实现Converter(可选)
如果您希望使用max-autotune模式功能,例如SuperKernel等高阶能力,需要额外实现Ascend Converter。Ascend Converter是指使用Ascend IR表达算子计算逻辑,在Eager模式下my_op调用Ascend C算子MyOp,对应到Converter实现,即调用Ascend IR MyOp。
- 在Ascend C算子工程编译时,除了生成aclnnXxx接口外,还会同步生成同名Ascend IR的注册代码。
- 接口介绍参见register_fx_node_ge_converter和custom_op。
在third_party/torchair/torchair/python/torchair/_ge_concrete_graph/ge_converter/custom目录下,新建MyOp算子对应的my_op.py文件,添加如下代码实现Converter:
1 2 3 4 5 6 7 | import torch import torchair # torch.ops.npu.my_op.default为自定义算子生成的Python函数签名,注意default后缀 @torchair.register_fx_node_ge_converter(torch.ops.npu.my_op.default) def convert_npu_my_op(x, y, z, attr1, attr2): # 函数入参与Torch算子入参一致 return torchair.ge.custom_op("MyOp", x, y, z, attr1, attr2) |
如果自定义算子同时满足如下条件,可跳过本章操作。
- 存在对应的Ascend IR,且Ascend IR的输入/输出顺序与PyTorch算子完全一致。
- TorchAir查找PyTorch算子对应的AscendIR时,按照如下规则查找:
- 先根据Python函数签名,获取算子名如torch.ops.npu.my_op.default中的my_op。
- 如果算子名中存在_functional后缀,会自动识别前面的算子名,如torch.ops.npu.my_inplace_functional.default中的my_inplace。
- 查找Ascend IR算子名时会根据找到的算子名,采用大驼峰命名法查找,如MyOp、MyInplace。
- PyTorch算子原型定义中存在Scalar类型入参时,不支持自动转换。
父主题: 非In-place算子开发和入图样例