5.8. Custom operations

PopRT supports creation of custom operators. You will need to create a custom operation when your model contains an operator that is not supported in PopRT. In this case, you can write a custom operator and compile it into a dynamic link library. PopRT supports dynamic linking of this custom operator into PopRT with the command line.

Since PopRT uses PopART as the backend, the process of developing custom operators for PopRT is the same as that for PopART. Please refer to Creating Custom Op in PopART .

This section describes the process of developing custom operators for PopRT with an example.

5.8.1. Writing custom operators

Taking the custom operator named LeakyRelu as an example, you first need to write the C++ code for it:

Listing 5.8 leaky_relu_custom_op.cpp

// Copyright (c) 2020 Graphcore Ltd. All rights reserved.

// This example demonstrates how to create a custom operator for PopART, in this
// case a Leaky ReLU op that returns `x` for any element `x >= 0` and `x *
// alpha` for any element `x < 0`, where `alpha` is provided as a scalar
// attribute to the operator.
#include <popart/operatoridentifier.hpp>
#include <popart/opmanager.hpp>
#include <popart/opserialiser.hpp>
#include <popart/popx/opxmanager.hpp>

#include <popops/ElementWise.hpp>
#include <popart/popx/opx.hpp>

namespace CustomOperators {
const popart::OperatorIdentifier LeakyReluId = {popart::Domain::ai_graphcore,
                                                "LeakyRelu",
                                                1};
} // namespace CustomOperators

class LeakyReluOp;
class LeakyReluOpx;

class LeakyReluOp : public popart::Op {
public:
  LeakyReluOp(const popart::OperatorIdentifier &_opid,
              float _alpha,
              const popart::Op::Settings &settings_)
      : popart::Op(_opid, settings_), alpha(_alpha) {}

  std::unique_ptr<Op> clone() const final {
    return std::make_unique<LeakyReluOp>(*this);
  }

  void setup() final { outInfo(0) = inInfo(0); }

  void appendAttributes(popart::OpSerialiserBase &os) const override {
    Op::appendAttributes(os);
    os.appendAttribute("alpha", getAlpha());
  }

  void appendOutlineAttributes(popart::OpSerialiserBase &os) const override {
    Op::appendOutlineAttributes(os);
    os.appendAttribute("alpha", getAlpha());
  }

  float getSubgraphValue() const final { return getHighSubgraphValue(); }

  bool requiresRandomSeed() const override { return false; }

  // Attributes
  float getAlpha() const { return alpha; }

private:
  float alpha;
};

namespace {
using popart::DataType;
using popart::OpDefinition;

static OpDefinition::DataTypes T = {DataType::FLOAT16, DataType::FLOAT};

static OpDefinition
    leakyReluOpDef({OpDefinition::Inputs({{"input", T}}),
                    OpDefinition::Outputs({{"output", T}}),
                    OpDefinition::Attributes({{"alpha", {"*"}}})});

static popart::OpCreator<LeakyReluOp> leakyReluOpCreator(
    popart::OpDefinitions({{CustomOperators::LeakyReluId, leakyReluOpDef}}),
    [](const popart::OpCreatorInfo &info) {
      // default alpha is 10**(-2)
      float alpha = info.attributes.getAttribute<popart::Attributes::Float>(
          "alpha", 1e-2f);
      return std::make_unique<LeakyReluOp>(info.opid, alpha, info.settings);
    },
    true);
} // namespace

namespace pe = popops::expr;

class LeakyReluOpx : public popart::popx::Opx {
public:
  LeakyReluOpx(popart::Op *op, popart::popx::Devicex *devicex)
      : popart::popx::Opx(op, devicex) {
    verifyOp<LeakyReluOp>(op, {CustomOperators::LeakyReluId});
  }

  void grow(poplar::program::Sequence &prog) const final {

    auto op = getOp<LeakyReluOp>();

    poplar::Tensor input = getInTensor(0);

    float alpha = op.getAlpha();

    // x < 0.0f ? alpha * x : x
    auto expression = pe::Select(pe::Mul(pe::Const(alpha), pe::_1),
                                 pe::_1,
                                 pe::Lt(pe::_1, pe::Const(0.0f)));

    popops::mapInPlace(graph(),
                       expression,
                       {input},
                       prog,
                       debugContext("LeakyRelu"),
                       poplar::OptionFlags());

    setOutTensor(0, input);
  }
};

static popart::popx::OpxCreator<LeakyReluOpx>
    LeakyReluOpxCreator({CustomOperators::LeakyReluId});

Download leaky_relu_custom_op.cpp

Create a Makefile and generate custom_ops.so using the make command:

Listing 5.9 Makefile

CXX ?= g++
CXXFLAGS = -std=c++14 -fPIC -g
LDLIBS = -shared -lpopart
ONNX_NAMESPACE = -DONNX_NAMESPACE=onnx

BUILD_DIR = build
SOURCES = leaky_relu_custom_op.cpp
TARGET = $(BUILD_DIR)/custom_ops.so

all: create_build_dir leaky_relu_custom_op

.PHONY: create_build_dir
create_build_dir:
	mkdir -p $(BUILD_DIR)

leaky_relu_custom_op: leaky_relu_custom_op.cpp
	$(CXX) $(SOURCES)  $(LDLIBS) $(CXXFLAGS) $(ONNX_NAMESPACE) -o $(TARGET)

.PHONY: clean
clean:
	rm -rf  $(BUILD_DIR)

Download Makefile

Create a shape-inference file for the custom operator:

Listing 5.10 custom_shape_inference.py

# Copyright (c) 2022 Graphcore Ltd. All rights reserved.
from typing import Tuple

import onnx
import onnx.helper
import onnx.shape_inference

from poprt.passes import ShapeFunc, get_dtype, get_shape, register_shape_func


@register_shape_func(['LeakyRelu'])
class LeakyRelu(ShapeFunc):
    """Function based on ONNX to infer the shape and dtype of custom op."""

    def __init__(self) -> None:
        super().__init__()

    def __call__(
        self,
        model: onnx.ModelProto,
        node: onnx.NodeProto,
    ) -> Tuple[onnx.ModelProto, bool]:
        graph = model.graph
        input_name = node.input[0]
        output_name = node.output[0]
        # If the Op already has known shape and dtype of output, return True
        if get_shape(model.graph, output_name) and get_dtype(model.graph, output_name):
            return model, True

        input_dtype = get_dtype(graph, input_name)
        input_shape = get_shape(graph, input_name)
        # If the Op is able to be inferred shape and dtype, return True
        if input_dtype and input_shape and 0 not in input_shape:
            # ![Shape-Inference Function begin]

            # Step.1: Write the method following ONNX-Protobuf standard,
            #         to calc shape and dtype of output in terms of shape and dtype of input
            # The LeakyRelu Op has same shape and dtype with input and output

            # Step.2: Create new TensorProto with inferred shape and dtype of output
            output_tensor = onnx.helper.make_tensor_value_info(
                output_name, input_dtype, input_shape
            )
            # Step.3: Call update_value_info to update
            model = self.update_value_info(model, output_tensor)
            # Step.4: Call infer_shapes function
            model = onnx.shape_inference.infer_shapes(model)
            # ![Shape-Inference Function end]
            return model, True
        # If the Op is not able to be inferred, return False
        else:
            return model, False

Download custom_shape_inference.py

Create the ONNX model file with the LeakyRelu op

Run the following test code with Python3 to generate the ONNX model file custom_op_test.onnx for testing:

Listing 5.11 create_onnx_with_custom_op.py

# Copyright (c) 2022 Graphcore Ltd. All rights reserved.
import argparse
import os

import onnx

from onnx import helper


def create_onnx_model_with_custom_op():
    TensorProto = onnx.TensorProto

    attributes = {"alpha": 0.01}
    leaky_relu = helper.make_node(
        "LeakyRelu", ["X"], ["Y"], domain="ai.graphcore", **attributes
    )
    relu = helper.make_node("Relu", ["Y"], ["Z"])

    graph = helper.make_graph(
        [leaky_relu, relu],
        "custom_op_test",
        [
            helper.make_tensor_value_info("X", TensorProto.FLOAT, (8, 8)),
        ],
        [
            helper.make_tensor_value_info("Z", TensorProto.FLOAT, (8, 8)),
        ],
    )
    opset_imports = [helper.make_opsetid("", 11)]
    model = helper.make_model(graph, opset_imports=opset_imports)
    model.opset_import.append(onnx.helper.make_opsetid("ai.graphcore", 1))
    return model


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Convert ONNX model and run it on the IPU.'
    )
    parser.add_argument(
        '--output_dir',
        type=str,
        default='./',
        help="Full path of the directory the ONNX model will be saved to.",
    )
    args = parser.parse_args()

    if not os.path.isdir(args.output_dir):
        raise ValueError("--output_dir should be an existing folder")

    model_path = os.path.join(args.output_dir, 'custom_op_test.onnx')

    model = create_onnx_model_with_custom_op()
    onnx.save(model, model_path)

    # Convert and Run
    compile_cmd = "bash build.sh"
    os.system(compile_cmd)
    abs_path = os.path.abspath(os.path.dirname(__file__))
    run_cmd = rf"""poprt \
--input_model {model_path} \
--custom_shape_inference {abs_path}/custom_shape_inference.py \
--custom_library_so_paths {abs_path}/custom_ops.so \
--run"""
    os.system(run_cmd)
    # 2022-12-30 07:01:54,408 INFO cli.py:446] Bs: 8
    # 2022-12-30 07:01:54,408 INFO cli.py:449] Latency: 0.23ms
    # 2022-12-30 07:01:54,408 INFO cli.py:450] Tput: 35469

Download create_onnx_with_custom_op.py

5.8.2. Using custom operators in PopRT

The library files of custom operators can be dynamically linked with the PopRT command line option --custom_library_so_paths, and the shape-inference for the custom operator can be registered with --custom_shape_inference.

The ONNX model file generated above can be executed with the following command:

poprt \
    --input_model custom_op_test.onnx \
    --custom_library_so_paths custom_ops.so \
    --custom_shape_inference custom_shape_inference.py \
    --run