1 引言：OMG插件是什么
OMG（Offline Model Generator，离线模型生成器）是华为昇腾AI处理器配套的模型转换工具。它将来自Caffe、TensorFlow、ONNX等框架的模型，转换成昇腾芯片能直接执行的 .om离线模型文件。
然而，深度学习领域发展迅速，新的算子层出不穷。OMG内置的算子库无法覆盖所有情况。这时就需要OMG自定义算子插件机制——它允许开发者编写昇腾原生不支持的算子，并将它们无缝集成到OMG的转换链路中。
本文将从原理到实践，完整地解析OMG插件的开发、使用和推理过程。
2 OMG插件原理：三段式架构
OMG插件的实现本质是一个“编译期扩展”框架。整个生命周期分为三个紧密相连的阶段：
flowchart LR
subgraph A [阶段一：开发阶段]
A1[编写Host侧代码<br>C++] --> A3[执行build.sh<br>编译工程]
A2[编写Kernel侧代码<br>Ascend C] --> A3
end
subgraph B [阶段二：OMG转换阶段]
B1[原始模型文件] --> B2[OMG加载]
B3[插件.so文件] --plugin_path--> B2
B2 --> B4[生成.o二进制]
B4 --> B5[打包成.om离线模型]
end
subgraph C [阶段三：OME推理阶段]
C1[输入数据] --> C2[OME加载.om]
C2 --> C3[AI Core/AI CPU执行]
C3 --> C4[输出结果]
end
A3 --> B3
B5 --> C2
3 阶段一：算子开发与插件构建
3.1 算子工程的组成部分
一个完整的自定义算子工程包含两个核心部分，在同一个工程目录下通过 msopgen工具生成框架：
| 组成部分 | 目录位置 | 开发语言 | 核心职责 | 编译产物 |
|---|---|---|---|---|
| Host侧 | op_host/ |
C++ | 算子原型注册、形状推导、Tiling策略 | .so插件库 |
| Kernel侧 | op_kernel/ |
Ascend C | AI Core上的核函数，核心计算逻辑 | .o二进制 + .json描述 |
说明：如果开发的是AI CPU算子，Kernel侧代码放在
cpukernel/impl/目录下，编译产物为.so动态库。
3.2 Host侧：算子的“说明书”
Host侧代码让OMG能够理解自定义算子：
- 算子原型注册：通过
REGISTER_OP宏将算子类型名与实现绑定，定义输入/输出个数、数据类型、属性参数。 - 形状推导（InferShape）：根据输入张量的形状计算出输出形状，帮助OMG做静态内存分配。
- Tiling策略计算：决定如何将大数据切分成小块，适配AI Core有限的片上存储。
3.3 Kernel侧：算子的“发动机”
Kernel侧代码运行在AI Core上，使用Ascend C编程范式开发：
// 一个典型的Ascend C核函数结构
extern "C" __global__ __aicore__
void add_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z)
{
KernelAdd op;
op.Init(x, y, z); // 初始化：获取块索引、切分数据
op.Process(); // 执行：CopyIn → Compute → CopyOut流水线
}
核心职责包括：
- 数据搬运（Global Memory ↔ Local Memory）
- 向量/矩阵计算指令调用
- 流水线并行调度（CopyIn/Compute/CopyOut三段式）
3.4 编译构建：build.sh的工作
执行 ./build.sh后，脚本自动完成：
- 编译Host侧代码：生成
libcust_opapi.so（调用入口）和libcust_op_proto.so（原型库） - 编译Kernel侧代码：使用Ascend C编译器将
.cpp核函数编译成.o和.json - 打包部署：将所有产物复制到
tools_omg目录，供OMG使用
编译支持两种模式：
- 二进制发布：预编译Kernel侧生成
.o，OMG直接使用 - 源码发布：保留
.cpp源码，支持OMG在线编译
4 阶段二：OMG模型转换
4.1 插件加载参数
OMG通过 --plugin_path参数指定自定义算子插件路径：
omg --model=my_net.prototxt \
--weight=my_net.caffemodel \
--framework=0 \
--output=my_net \
--plugin_path=/path/to/plugin # 指向包含.so的目录
4.2 转换流程
flowchart TD
A[OMG解析原始模型] --> B{遇到算子节点}
B --> C[在插件.so中查找算子原型]
C --> D[获取输入输出描述]
D --> E[调用InferShape推导输出形状]
E --> F{算子执行部件类型}
F -->|ImplyType::TVM| G[AI Core路径]
F -->|ImplyType::AI_CPU| H[AI CPU路径]
G --> I[调用TBE/Ascend C编译器]
I --> J[生成.o二进制]
H --> K[调用C++编译器]
K --> L[生成.so动态库]
J --> M[打包.om离线模型]
L --> M
4.3 执行部件的区分
通过Host侧注册时的 ImplyType参数，OMG决定算子运行的硬件单元：
| ImplyType值 | 执行硬件 | 适用算子类型 |
|---|---|---|
ImplyType::TVM |
AI Core | 计算密集型（卷积、矩阵乘、Add等） |
ImplyType::AI_CPU |
AI CPU | 控制密集型（NMS、TopK、Reshape等） |
4.4 最终产出：.om离线模型
OMG将所有编译产物（.o、.so、.json）连同模型结构信息打包成一个独立的 .om离线模型文件，该文件包含昇腾AI处理器执行推理所需的全部指令和数据。
5 阶段三：OME推理执行
5.1 执行架构
离线模型执行器（OME）在推理时加载和执行 .om模型：
flowchart LR
A[应用程序] --> B[模型管家]
B --> C[加载.om离线模型]
C --> D[解析模型中的算子]
D --> E{算子类型}
E -->|AI Core算子| F[任务调度器]
F --> G[AI Core执行.o指令]
E -->|AI CPU算子| H[AI CPU Engine]
H --> I[AI CPU执行.so逻辑]
G --> J[输出结果]
I --> J
5.2 AI Core算子的执行细节
- 多核并行：输入数据按
GetBlockIdx()分配到不同AI Core并行处理 - Tile切分：每个Core内部将数据切分成小块，适配Local Memory
- 双缓冲流水线：CopyIn、Compute、CopyOut三阶段重叠执行，最大化利用率
5.3 AI CPU算子的执行细节
对于 ImplyType::AI_CPU算子：
- OME通过
dlopen()加载对应的.so动态库 - 调用其中注册的
Compute函数在AI CPU上执行 - 适用于逻辑复杂、分支多、难以向量化的操作
6 开发选型建议
6.1 算子开发路线对比
| 对比维度 | TBE (Python) | Ascend C (C++) |
|---|---|---|
| 当前状态 | 即将废弃 | 官方力推 |
| 开发语言 | Python | C++ |
| 性能上限 | 较高 | 最高 |
| 适合场景 | 快速验证、存量维护 | 所有新开发项目 |
结论：新开发统一使用Ascend C。
6.2 AI Core vs AI CPU 选择指南
| 算子类型 | 推荐执行部件 | 原因 |
|---|---|---|
| 卷积、矩阵乘、池化 | AI Core | 计算密集型，硬件高度优化 |
| Add、Mul、ReLU | AI Core | 向量化友好，效率高 |
| NMS、Where、TopK | AI CPU | 逻辑分支多，难以向量化 |
| 动态形状算子 | AI CPU | AI Core对动态形状支持有限 |
7 总结：OMG插件的完整生命周期
flowchart TD
subgraph Dev[开发阶段]
D1[编写Host侧C++代码] --> D3[执行build.sh]
D2[编写Kernel侧Ascend C] --> D3
D3 --> D4[生成.so插件库]
D3 --> D5[生成.o二进制]
end
subgraph OMG[OMG转换阶段]
O1[原始模型] --> O2[OMG加载]
O3[.so插件] --plugin_path--> O2
O2 --> O4[解析原型、形状推导]
O4 --> O5[调用编译器生成代码]
O5 --> O6[打包.om离线模型]
end
subgraph OME[OME推理阶段]
M1[输入数据] --> M2[OME加载.om]
M2 --> M3[解析任务流]
M3 --> M4{算子类型}
M4 -->|AI Core| M5[调度器→AI Core执行.o]
M4 -->|AI CPU| M6[AI CPU执行.so]
M5 --> M7[输出]
M6 --> M7
end
D4 --> O3
D5 --> O5
O6 --> M2
核心要点：
- 开发阶段：用Ascend C编写Host侧和Kernel侧代码，
build.sh一次性编译生成.so和.o - 转换阶段：OMG通过
--plugin_path加载.so插件，将自定义算子编译并打包进.om模型 - 推理阶段：OME根据模型中记录的执行部件类型，AI Core算子跑在AI Core执行
.o指令，AI CPU算子跑在AI CPU执行.so逻辑
通过这套机制，开发者可以将任意自定义算子无缝集成到昇腾推理链路中，充分发挥昇腾AI处理器的计算潜力。