ONNX Plugin Python 接口方案评审

文档说明

本文用于评审 GE ONNX Plugin 的 Python 公开接口。文档首先说明现有 C++ ONNX 到 GE IR 的链路，再比较 Python Plugin 的候选接口，最后给出当前推荐的用户界面。

本次评审范围包括：

• Python Plugin 组织方式；
• Python Plugin 发现方式

1. GE C++ plugin组织形式

当前 C++ Plugin 的典型写法如下：

REGISTER_CUSTOM_OP("Elu")
    .FrameworkType(domi::ONNX)
    .OriginOpType({
        "ai.onnx::8::Elu",
        "ai.onnx::9::Elu",
        "ai.onnx::10::Elu",
        "ai.onnx::11::Elu"
    })
    .ImplyType(domi::ImplyType::TVM)
    .ParseParamsFn(ParseEluParams);

注册字段及其消费者如下：

C++ 注册字段	保存内容	ONNX 链路消费者	作用
FrameworkType(ONNX)	framework 枚举	初始化入口、OpParserFactory	选择 ONNX 注册数据和 parser factory。
OriginOpType	完整 source type 集合	OpRegistry::Register、AdapterOpType	建立 source 到 target 的映射并参与 callback 查找。
REGISTER_CUSTOM_OP("Elu")	GE target type	Finalize、OpParserFactory、OperatorFactory	注册 parser creator并创建 target Operator。
ParseParamsFn	parser callback	OnnxCustomParserAdapter	读取 source 节点并补充 target Operator。
ImplyType	声明性运行模式	OpRegistry::Register、初始化日志	写入 op_run_mode_map_；ONNX parser 不读取该值。

2. 本次评审细节

2.2 Python Plugin 组织方式

2.2.1 方案一：单装饰器

@onnx_plugin(
    source="Elu",
    domain="ai.onnx",
    opsets=range(8, 19),
    target="Elu",
)
def parse_elu(node: OnnxNode, target: Operator) -> None:
    target.set_attr("alpha", node.attrs.get("alpha", 1.0))

优点是首轮代码最短。主要问题是 descriptor 与一种 callback 直接绑定。后续增加一对多分解、ByOperator 或 subgraph callback 时，需要增加新的顶层装饰器或重复填写 source/domain/opsets/target。
a

2.2.2 方案二：descriptor + callback

elu = onnx_plugin(
    source="Elu",
    domain="ai.onnx",
    opsets=range(8, 19),
    target="Elu",
)


@elu.parse_node
def parse_elu(node: OnnxNode, target: Operator) -> None:
    target.set_attr("alpha", node.attrs.get("alpha", 1.0))


# 一对多场景
@elu.decompose
def elu_decompose(source, builder):
    ...

descriptor 集中保存稳定的 source 到 target 描述，不同 callback 作为独立成员扩展。首轮只公开 parse_node，未来增加 decompose 等能力时不需要改变 descriptor 构造方式，也不需要重复注册字段。

2.2.3 方案三：class

@onnx_plugin(
    source="Elu",
    domain="ai.onnx",
    opsets=range(8, 19),
    target="Elu",
)
class EluPlugin:
    def parse_node(self, node: OnnxNode, target: Operator) -> None:
        target.set_attr("alpha", node.attrs.get("alpha", 1.0))

    def decompose(self, source, builder):
        ...

class 可以聚合多个方法，但会为无状态 parser callback 引入实例、self 和对象生命周期，用户还需要理解框架何时创建 Plugin 实例。

2.2.4 方案对比

方案	首轮易用性	后续 callback 扩展	状态和生命周期	结论
单装饰器	最短	容易重复字段或增加顶层装饰器	无额外状态	不推荐
descriptor + callback	多一个 descriptor 变量	descriptor 不变，增加成员即可	descriptor 进程级，callback 普通函数	推荐
class	方法聚合	可以增加方法	引入不必要实例状态	不推荐

2.2.5 ImplyType 处理结论

源码核验确认注册 ImplyType 在 ONNX 链路中没有功能性消费者。Python descriptor 不提供 imply_type 参数，模块不导出 ImplyType Enum，内部保持 OpRegistrationData 默认 BUILDIN。该值不影响 ONNX 解析、target 创建和引擎选择。

后续只有在出现明确 ONNX 消费者和用户语义时，才通过独立公开接口评审新增该字段。

2.3 Python Plugin 发现方式

2.3.1 方案一：新增 ONNX 专用变量

export ASCEND_ONNX_PLUGIN_PATH=/path/to/plugin.py

该方案语义独立，但增加产品级环境变量、部署入口和文档。一个文件同时声明 Python custom_op 和 ONNX Plugin 时还需要重复配置。

2.3.2 方案二：复用现有直接扫描

export ASCEND_CUSTOM_OPP_PATH=/path/to/plugin.py

该方案改动最小，但会继续混合 OPP root、单 Python 文件、普通目录和 package 语义，普通 OPP root 下的无关 Python 文件可能被扫描。

推荐方案二

3. 推荐用户界面

3.1 Elu 示例

from ge.graph import Operator
from ge.onnx_plugin import OnnxNode, onnx_plugin


elu = onnx_plugin(
    source="Elu",
    domain="ai.onnx",
    opsets=range(8, 19),
    target="Elu",
)


@elu.parse_node
def parse_elu(node: OnnxNode, target: Operator) -> None:
    alpha = node.attrs.get("alpha", 1.0)
    if not isinstance(alpha, float):
        raise TypeError("Elu alpha must be float")
    target.set_attr("alpha", alpha)

用户只读取 source 属性并修改已有 target

3.2 后续 callback 扩展

descriptor + callback 允许未来增加独立转换阶段：

# 示意接口，名称、参数和返回值不属于首轮承诺。
@elu.decompose
def elu_decompose(source, builder):
    ...

首轮 OnnxPlugin 只公开 parse_node，不提前加入 decompose 占位方法。

4竞品用户界面

4.1 PyTorch ONNX Exporter

program = torch.onnx.export(
    model,
    args,
    dynamo=True,
    custom_translation_table={
        torch.ops.aten.add.Tensor: custom_add,
    },
)

其方向是 PyTorch/FX 到 ONNX，与本需求方向相反。可有限参考普通 Python callable 和诊断方式。

4.2 TensorRT ONNX Parser

parser = trt.OnnxParser(network, logger)
if not parser.parse(model_bytes):
    for index in range(parser.num_errors):
        print(parser.get_error(index))

TensorRT 的 ONNX 到 Network 映射位于内部 C++ importer，Python 用户不能注册新的节点映射。可参考 Parser/Network 分层和节点级错误。

4.3 TensorRT Python Plugin

@trtp.register("example::identity")
def identity_desc(inp: trtp.TensorDesc) -> trtp.TensorDesc:
    return inp.like()


@trtp.impl("example::identity")
def identity_impl(inp, outputs, stream) -> None:
    launch_identity(inp, outputs[0], stream)

该接口定义自定义层的 shape、dtype 和运行时 Kernel，不是 ONNX source 到已有 GE Operator 的映射。

4.4 onnx2torch

模型用户入口：

from onnx2torch import convert

torch_model = convert("model.onnx")

converter 开发者入口：

@add_converter(operation_type="Elu", version=6)
def convert_elu(node: OnnxNode, graph: OnnxGraph) -> OperationConverterResult:
    return OperationConverterResult(
        torch_module=nn.ELU(alpha=node.attributes.get("alpha", 1.0)),
        onnx_mapping=onnx_mapping_from_node(node),
    )

onnx2torch 从零创建 PyTorch Module 和 FX Graph，因此需要 OnnxGraph、OperationConverterResult 和 OnnxMapping。GE parser 已经创建 target Operator并负责入图和连边，不需要照搬这些对象。

4.5 横向结论

对象	转换方向	用户 callback 的主要职责	对本方案的参考
PyTorch ONNX Exporter	PyTorch/FX -> ONNX	创建 ONNX 表达	callable 和诊断
TensorRT ONNX Parser	ONNX -> TensorRT Network	Python 不可扩展节点映射	Parser 分层和错误
TensorRT Python Plugin	Custom Node -> Plugin Layer	shape/dtype/Kernel	装饰器形式
onnx2torch	ONNX -> PyTorch FX/Module	创建 Module 和 mapping	只读节点和属性默认值写法
GE Python ONNX Plugin	ONNX -> 已有 GE target Operator	补充属性和实例端口	本次设计主体