paddlepaddle--paddle/test/ipu/custom_ops

Add custom op for Paddle on IPU

Add custom op in Paddle

reference

https://www.paddlepaddle.org.cn/documentation/docs/zh/guides/custom_op/new_cpp_op_cn.html

Write custom op for PopART

reference

https://docs.graphcore.ai/projects/popart-user-guide/en/latest/custom_ops.html

Register custom op for Paddle on IPU

这里采用即时编译(JIT Compile) 的方法使用 custom op.

实现 custom op

根据上面的两个文档, 首先添加 custom op 的实现.

leaky_relu_cpu.cc 包含了 Paddle 中 custom op 的定义和 cpu 实现, 这里的实现是和标准的 Paddle 添加 custom op 是完全一致的. 这里的 cpu 实现不是必须的, cpu 实现可以用来检验 ipu 实现的正确性.

leaky_relu_ipu.cc 包含了 PopART 中 custom op 的定义和 ipu 实现, 同样的, 这里的实现和标准的 PopART 添加 custom op 是完全一致的.

载入 custom op

分别在 Paddle 和 PopART 中实现 custom op 的定义后, 使用 paddle.utils.cpp_extension.load 编译源文件并把对应的动态库加载到当前进程中.

cur_dir = os.path.dirname(os.path.realpath(__file__))
custom_ops = load(
    name="custom_jit_ops",
    sources=[
        f"{cur_dir}/leaky_relu_cpu.cc",
        f"{cur_dir}/leaky_relu_ipu.cc",
    ],
    # 编译 leaky_relu_ipu.cc 时需要添加此参数
    extra_cxx_cflags=['-DONNX_NAMESPACE=onnx'])

由于 Paddle 中 op 的定义和 PopART 中存在一些差异, 需要手动映射 custom op

# paddle_op is custom op type in Paddle
# popart_op, domain and version is custom op identifier in PopART
ipu_strategy = paddle.static.IpuStrategy()
ipu_strategy.add_custom_op(
    paddle_op="custom_leaky_relu",
    popart_op="LeakyRelu",
    domain='custom.ops',
    version=1)

使用 custom op

x = paddle.static.data(
    name=self.feed_list[0],
    shape=self.feed_shape[0],
    dtype=self.feed_dtype[0])
# custom op
out = custom_ops.custom_leaky_relu(x, **self.attrs)