5.6. 开发 Custom Operation

PopRT 支持用户开发自定义算子, 用于对 PopRT 做扩展.

典型的应用场景是: 用户有一个 ONNX 模型, 其中某个算子在 PopRT 中不支持, 此时用户就可以编写一个自定义算子, 并编译成动态链接库, PopRT 支持通过命令行的方式把这个自定义算子动态链接进 PopRT.

下面通过一个例子来描述为 PopRT 开发自定义算子的流程.

5.6.1. 编写自定义算子

由于 PopRT 是用 PopART 作为 backend, 因此为 PopRT 开发自定义算子的流程和 PopART 一致, 请参考 Creating Custom OP in PopART.

以名为 LeakyRelu 的自定义算子为例, 首先需要编写一个自定义算子的 C++ 代码:

Listing 5.8 leaky_relu_custom_op.cpp
  1// Copyright (c) 2020 Graphcore Ltd. All rights reserved.
  2
  3// This example demonstrates how to create a custom operator for PopART, in this
  4// case a Leaky ReLU op that returns `x` for any element `x >= 0` and `x *
  5// alpha` for any element `x < 0`, where `alpha` is provided as a scalar
  6// attribute to the operator.
  7#include <popart/operatoridentifier.hpp>
  8#include <popart/opmanager.hpp>
  9#include <popart/opserialiser.hpp>
 10#include <popart/popx/opxmanager.hpp>
 11
 12#include <popops/ElementWise.hpp>
 13#include <popart/popx/opx.hpp>
 14
 15namespace CustomOperators {
 16const popart::OperatorIdentifier LeakyReluId = {popart::Domain::ai_graphcore,
 17                                                "LeakyRelu",
 18                                                1};
 19} // namespace CustomOperators
 20
 21class LeakyReluOp;
 22class LeakyReluOpx;
 23
 24class LeakyReluOp : public popart::Op {
 25public:
 26  LeakyReluOp(const popart::OperatorIdentifier &_opid,
 27              float _alpha,
 28              const popart::Op::Settings &settings_)
 29      : popart::Op(_opid, settings_), alpha(_alpha) {}
 30
 31  std::unique_ptr<Op> clone() const final {
 32    return std::make_unique<LeakyReluOp>(*this);
 33  }
 34
 35  void setup() final { outInfo(0) = inInfo(0); }
 36
 37  void appendAttributes(popart::OpSerialiserBase &os) const override {
 38    Op::appendAttributes(os);
 39    os.appendAttribute("alpha", getAlpha());
 40  }
 41
 42  void appendOutlineAttributes(popart::OpSerialiserBase &os) const override {
 43    Op::appendOutlineAttributes(os);
 44    os.appendAttribute("alpha", getAlpha());
 45  }
 46
 47  float getSubgraphValue() const final { return getHighSubgraphValue(); }
 48
 49  bool requiresRandomSeed() const override { return false; }
 50
 51  // Attributes
 52  float getAlpha() const { return alpha; }
 53
 54private:
 55  float alpha;
 56};
 57
 58namespace {
 59using popart::DataType;
 60using popart::OpDefinition;
 61
 62static OpDefinition::DataTypes T = {DataType::FLOAT16, DataType::FLOAT};
 63
 64static OpDefinition
 65    leakyReluOpDef({OpDefinition::Inputs({{"input", T}}),
 66                    OpDefinition::Outputs({{"output", T}}),
 67                    OpDefinition::Attributes({{"alpha", {"*"}}})});
 68
 69static popart::OpCreator<LeakyReluOp> leakyReluOpCreator(
 70    popart::OpDefinitions({{CustomOperators::LeakyReluId, leakyReluOpDef}}),
 71    [](const popart::OpCreatorInfo &info) {
 72      // default alpha is 10**(-2)
 73      float alpha = info.attributes.getAttribute<popart::Attributes::Float>(
 74          "alpha", 1e-2f);
 75      return std::make_unique<LeakyReluOp>(info.opid, alpha, info.settings);
 76    },
 77    true);
 78} // namespace
 79
 80namespace pe = popops::expr;
 81
 82class LeakyReluOpx : public popart::popx::Opx {
 83public:
 84  LeakyReluOpx(popart::Op *op, popart::popx::Devicex *devicex)
 85      : popart::popx::Opx(op, devicex) {
 86    verifyOp<LeakyReluOp>(op, {CustomOperators::LeakyReluId});
 87  }
 88
 89  void grow(poplar::program::Sequence &prog) const final {
 90
 91    auto op = getOp<LeakyReluOp>();
 92
 93    poplar::Tensor input = getInTensor(0);
 94
 95    float alpha = op.getAlpha();
 96
 97    // x < 0.0f ? alpha * x : x
 98    auto expression = pe::Select(pe::Mul(pe::Const(alpha), pe::_1),
 99                                 pe::_1,
100                                 pe::Lt(pe::_1, pe::Const(0.0f)));
101
102    popops::mapInPlace(graph(),
103                       expression,
104                       {input},
105                       prog,
106                       debugContext("LeakyRelu"),
107                       poplar::OptionFlags());
108
109    setOutTensor(0, input);
110  }
111};
112
113static popart::popx::OpxCreator<LeakyReluOpx>
114    LeakyReluOpxCreator({CustomOperators::LeakyReluId});

Download leaky_relu_custom_op.cpp

编写 Makefile 并通过 make 命令生成 custom_ops.so:

Listing 5.9 Makefile
 1CXX ?= g++
 2CXXFLAGS = -std=c++14 -fPIC -g
 3LDLIBS = -shared -lpopart
 4ONNX_NAMESPACE = -DONNX_NAMESPACE=onnx
 5
 6BUILD_DIR = build
 7SOURCES = leaky_relu_custom_op.cpp
 8TARGET = $(BUILD_DIR)/custom_ops.so
 9
10all: create_build_dir leaky_relu_custom_op
11
12.PHONY: create_build_dir
13create_build_dir:
14	mkdir -p $(BUILD_DIR)
15
16leaky_relu_custom_op: leaky_relu_custom_op.cpp
17	$(CXX) $(SOURCES)  $(LDLIBS) $(CXXFLAGS) $(ONNX_NAMESPACE) -o $(TARGET)
18
19.PHONY: clean
20clean:
21	rm -rf  $(BUILD_DIR)

Download Makefile

编写自定义算子的 Shape-Inference 文件:

Listing 5.10 custom_shape_inference.py
 1# Copyright (c) 2022 Graphcore Ltd. All rights reserved.
 2from typing import Tuple
 3
 4import onnx
 5import onnx.helper
 6import onnx.shape_inference
 7
 8from poprt.passes import ShapeFunc, get_dtype, get_shape, register_shape_func
 9
10
11@register_shape_func(['LeakyRelu'])
12class LeakyRelu(ShapeFunc):
13    """Function based on ONNX to infer the shape and dtype of Custom Op."""
14
15    def __init__(self) -> None:
16        super().__init__()
17
18    def __call__(
19        self,
20        model: onnx.ModelProto,
21        node: onnx.NodeProto,
22    ) -> Tuple[onnx.ModelProto, bool]:
23        graph = model.graph
24        input_name = node.input[0]
25        output_name = node.output[0]
26        # If the Op already has known shape and dtype of output, return True
27        if get_shape(model.graph, output_name) and get_dtype(model.graph, output_name):
28            return model, True
29
30        input_dtype = get_dtype(graph, input_name)
31        input_shape = get_shape(graph, input_name)
32        # If the Op is able to be inferred shape and dtype, return True
33        if input_dtype and input_shape and 0 not in input_shape:
34            # ![Shape-Inference Function begin]
35
36            # Step.1: Write the method following ONNX-Protobuf standard,
37            #         to calc shape and dtype of output in terms of shape and dtype of input
38            # The LeakyRelu Op has same shape and dtype with input and output
39
40            # Step.2: Create new TensorProto with inferred shape and dtype of output
41            output_tensor = onnx.helper.make_tensor_value_info(
42                output_name, input_dtype, input_shape
43            )
44            # Step.3: Call update_value_info to update
45            model = self.update_value_info(model, output_tensor)
46            # Step.4: Call infer_shapes function
47            model = onnx.shape_inference.infer_shapes(model)
48            # ![Shape-Inference Function end]
49            return model, True
50        # If the Op is not able to be inferred, return False
51        else:
52            return model, False

Download custom_shape_inference.py

创建一个带有 LeakyRelu OP 的 ONNX 模型文件

通过 Python3 运行以下测试代码生成用于测试的 ONNX 模型文件 custom_op_test.onnx:

Listing 5.11 create_onnx_with_custom_op.py
 1# Copyright (c) 2022 Graphcore Ltd. All rights reserved.
 2import argparse
 3import os
 4
 5import onnx
 6
 7from onnx import helper
 8
 9
10def create_onnx_model_with_custom_op():
11    TensorProto = onnx.TensorProto
12
13    attributes = {"alpha": 0.01}
14    leaky_relu = helper.make_node(
15        "LeakyRelu", ["X"], ["Y"], domain="ai.graphcore", **attributes
16    )
17    relu = helper.make_node("Relu", ["Y"], ["Z"])
18
19    graph = helper.make_graph(
20        [leaky_relu, relu],
21        "custom_op_test",
22        [
23            helper.make_tensor_value_info("X", TensorProto.FLOAT, (8, 8)),
24        ],
25        [
26            helper.make_tensor_value_info("Z", TensorProto.FLOAT, (8, 8)),
27        ],
28    )
29    opset_imports = [helper.make_opsetid("", 11)]
30    model = helper.make_model(graph, opset_imports=opset_imports)
31    model.opset_import.append(onnx.helper.make_opsetid("ai.graphcore", 1))
32    return model
33
34
35if __name__ == '__main__':
36    parser = argparse.ArgumentParser(
37        description='Convert onnx model and run it on IPU.'
38    )
39    parser.add_argument(
40        '--output_dir',
41        type=str,
42        default='./',
43        help="Full path of the onnx model will be saved to.",
44    )
45    args = parser.parse_args()
46
47    if not os.path.isdir(args.output_dir):
48        raise ValueError("--output_dir should be an exist folder")
49
50    model_path = os.path.join(args.output_dir, 'custom_op_test.onnx')
51
52    model = create_onnx_model_with_custom_op()
53    onnx.save(model, model_path)
54
55    # Convert and Run
56    compile_cmd = "bash build.sh"
57    os.system(compile_cmd)
58    abs_path = os.path.abspath(os.path.dirname(__file__))
59    run_cmd = rf"""poprt \
60--input_model {model_path} \
61--custom_shape_inference {abs_path}/custom_shape_inference.py \
62--custom_library_so_paths {abs_path}/custom_ops.so \
63--run"""
64    os.system(run_cmd)
65    # 2022-12-30 07:01:54,408 INFO cli.py:446] Bs: 8
66    # 2022-12-30 07:01:54,408 INFO cli.py:449] Latency: 0.23ms
67    # 2022-12-30 07:01:54,408 INFO cli.py:450] Tput: 35469

Download create_onnx_with_custom_op.py

在 PopRT 中使用自定义算子

可以通过 PopRT 的命令行 --custom_library_so_paths 来动态链接自定义算子的库文件, 并通过 --custom_shape_inference 来注册自定义算子的 Shape-Inference.

可以通过如下命令来执行上述生成的 ONNX 模型文件:

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