Custom nGraph Operation

Inference Engine Extension API allows to register operation sets (opsets) with custom nGraph operations, it allows to support Networks with unknown operations.

Operation Class

To add your custom nGraph operation, create a new class that extends ngraph::Op, which is in turn derived from ngraph::Node, the base class for all graph operations in nGraph. Follow the steps below:

  1. Define a NodeTypeInfo object that identifies the type of the operation to the graph users and helps with dynamic type resolution. The type info of an nGraph operation currently consists of a string identifier and a version number, but this may change in the future.
  2. Implement constructors that can optionally take the operation inputs and attributes as parameters.
  3. Override the shape inference method validate_and_infer_types. This method is called multiple times during graph manipulations to determine the shapes and element types of the outputs of the operations. You can access the input shapes through the get_input_partial_shape() method and input element types through the get_input_element_type() method of ngraph::Node. Set the inferred shape and element type of the output using set_output_type.
  4. Override the clone_with_new_inputs method, which allows graph manipulation routines to create copies of this operation and connect it to different nodes during optimization.
  5. Override the visit_attributes method, which allows serialization and deserialization of attributes. An AttributeVisitor is passed to the method, and the implementation is expected to walk over all the attributes in the op using the type-aware on_attribute helper. Helpers are already implemented for standard C++ types like int64_t, float, bool, vector and for existing nGraph defined types.
  6. Override evaluate, which is an optional method that enables the application of constant folding if there is a custom operation on the constant branch.

Based on that, declaration of a operation class can look as follows:

namespace TemplateExtension {
class Operation : public ngraph::op::Op {
public:
static constexpr ngraph::NodeTypeInfo type_info{"Template", 0};
const ngraph::NodeTypeInfo& get_type_info() const override { return type_info; }
Operation() = default;
Operation(const ngraph::Output<ngraph::Node>& arg, int64_t add);
void validate_and_infer_types() override;
std::shared_ptr<ngraph::Node> clone_with_new_inputs(const ngraph::OutputVector& new_args) const override;
bool visit_attributes(ngraph::AttributeVisitor& visitor) override;
int64_t getAddAttr() const { return add; }
bool evaluate(const ngraph::HostTensorVector& outputs,
const ngraph::HostTensorVector& inputs) const override;
private:
int64_t add;
};

Class Fields

The provided implementation has several fields:

  • add of type int64_t is an attribute of custom operation
  • type_info of type ngraph::NodeTypeInfo defines the type and version of operation

Operation Constructors

nGraph operation contains two constructors: a default constructor, which allows to create operation without attributes and a constructor that creates and validates operation with specified inputs and attributes.

Operation::Operation(const ngraph::Output<ngraph::Node> &arg, int64_t add) : Op({arg}), add(add) {
constructor_validate_and_infer_types();
}

validate_and_infer_types()

ngraph::Node::validate_and_infer_types method validates operation attributes and calculates output shapes using attributes of operation.

void Operation::validate_and_infer_types() {
// Operation doesn't change shapes end element type
set_output_type(0, get_input_element_type(0), get_input_partial_shape(0));
}

clone_with_new_inputs()

ngraph::Node::clone_with_new_inputs method creates a copy of nGraph operation with new inputs.

std::shared_ptr<ngraph::Node> Operation::clone_with_new_inputs(const ngraph::OutputVector &new_args) const {
if (new_args.size() != 1) {
throw ngraph::ngraph_error("Incorrect number of new arguments");
}
return std::make_shared<Operation>(new_args.at(0), add);
}

visit_attributes()

ngraph::Node::visit_attributes method allows to visit all operation attributes.

bool Operation::visit_attributes(ngraph::AttributeVisitor &visitor) {
visitor.on_attribute("add", add);
return true;
}

evaluate()

ngraph::Node::evaluate method allows to apply constant folding to an operation.

namespace
{
template <class T>
void implementation(const T* input,
T* output,
int64_t add,
size_t size) {
for (size_t i = 0; i < size; i++) {
output[i] = input[i] + add;
}
}
template <ngraph::element::Type_t ET>
bool evaluate_op(const ngraph::HostTensorPtr& arg0,
const ngraph::HostTensorPtr& out, int64_t add)
{
size_t size = ngraph::shape_size(arg0->get_shape());
implementation(arg0->get_data_ptr<ET>(),
out->get_data_ptr<ET>(),
add,
size);
return true;
}
} // namespace
bool Operation::evaluate(const ngraph::HostTensorVector& outputs,
const ngraph::HostTensorVector& inputs) const {
switch (inputs[0]->get_element_type())
{
case ngraph::element::Type_t::i8: return evaluate_op<ngraph::element::Type_t::i8>(inputs[0], outputs[0], getAddAttr());
case ngraph::element::Type_t::i16: return evaluate_op<ngraph::element::Type_t::i16>(inputs[0], outputs[0], getAddAttr());
case ngraph::element::Type_t::i32: return evaluate_op<ngraph::element::Type_t::i32>(inputs[0], outputs[0], getAddAttr());
case ngraph::element::Type_t::i64: return evaluate_op<ngraph::element::Type_t::i64>(inputs[0], outputs[0], getAddAttr());
case ngraph::element::Type_t::u8: return evaluate_op<ngraph::element::Type_t::u8>(inputs[0], outputs[0], getAddAttr());
case ngraph::element::Type_t::u16: return evaluate_op<ngraph::element::Type_t::u16>(inputs[0], outputs[0], getAddAttr());
case ngraph::element::Type_t::u32: return evaluate_op<ngraph::element::Type_t::u32>(inputs[0], outputs[0], getAddAttr());
case ngraph::element::Type_t::u64: return evaluate_op<ngraph::element::Type_t::u8>(inputs[0], outputs[0], getAddAttr());
case ngraph::element::Type_t::bf16: return evaluate_op<ngraph::element::Type_t::bf16>(inputs[0], outputs[0], getAddAttr());
case ngraph::element::Type_t::f16: return evaluate_op<ngraph::element::Type_t::f16>(inputs[0], outputs[0], getAddAttr());
case ngraph::element::Type_t::f32: return evaluate_op<ngraph::element::Type_t::f32>(inputs[0], outputs[0], getAddAttr());
default: break;
}
return false;
}

Register Custom Operations in Extension Class

To add custom operations to the Extension class, create an operation set with custom operations and implement the InferenceEngine::IExtension::getOpSets method:

std::map<std::string, ngraph::OpSet> Extension::getOpSets() {
std::map<std::string, ngraph::OpSet> opsets;
ngraph::OpSet opset;
opset.insert<Operation>();
opsets["custom_opset"] = opset;
return opsets;
}

This method returns a map of opsets that exist in the extension library.

nGraph provides opsets mechanism for operation versioning. Different opsets distinguish between different versions of one operation.

When specifying opset names, follow the rules below:

  • Use unique opset names.
  • Do not use the following built-in opset names: extension, experimental, opset1, opset2, opset3, ... , opsetN.
  • Make sure that the Model Optimizer and your extension use the same opset names.
  • IR v10 operations have the mandatory version attribute specifying the opset. Operations from the default opset cannot be redefined.

Use a custom opset to create a new operation or extend functionality of an existing operation from another opset.