编译执行链路

源代码：https://gitcode.com/cann 主要源于ops-nn; ops-math算子库; ge图引擎库; 运行时runtime库

执行过程

单算子

单算子链路：

PyTorch (用户代码)
  └─ Dispatcher（根据 tensor.device 选择 PrivateUse1 后端）
       └─ op-plugin / torch_npu（TORCH_LIBRARY_IMPL 注册的 NPU 实现）
            └─ aclnn 接口（libopapi_nn.so，CPU 侧调度层）
                 └─ CANN Runtime（执行 tiling，按 .json 索引匹配 kernel）
                      └─ AscendC Kernel（.o binary，在 NPU 上执行）

op-plugin（pytorch_adapter）和 torch_npu 属于同一层——op-plugin 以插件方式向 torch_npu 注册算子实现，对外统称 “NPU 后端”。

用户代码：

tensor = torch.randn(1024).npu()
out = torch.relu(tensor)

第一层：PyTorch Dispatcher

torch.relu(tensor)
  └─ tensor.device == npu
       └─ 选择 PrivateUse1 dispatch key
            └─ 调用 op-plugin 注册的实现
               （TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) { m.impl("relu", relu_npu) }）

第二层：op-plugin wrapper（CPU 侧）

at::Tensor relu_npu(const at::Tensor& self) {
    // 1. InferShape：推导输出 shape，分配输出 tensor
    at::Tensor result = npu_preparation::apply_tensor(self);
    // 2. 查询 workspace 大小（内部触发 tiling）
    uint64_t workspaceSize = 0;
    aclOpExecutor* executor = nullptr;
    aclnnReluGetWorkspaceSize(self, result, &workspaceSize, &executor);
    // 3. 分配 NPU workspace（临时显存）
    void* workspace = ctx.allocate(workspaceSize);
    // 4. 提交执行
    aclnnRelu(workspace, workspaceSize, executor, stream);
    return result;
}

第三层：aclnnReluGetWorkspaceSize()（libopapi_nn.so）

调用 libophost_ops_nn.so 中的 tiling 函数：

输入：实际 shape、dtype、硬件规格（UB 大小、核数等）
输出：TilingData（每个 AI Core 处理多少数据、分几块等）
输出：workspace 大小
返回：executor handle（携带 TilingData）

第四层：aclnnRelu()（libopapi_nn.so）

CANN Runtime：

查 build/binary/config/ascend950/*.json（运行时全局索引）
- 匹配条件：op_name=Relu + dtype=float32
- 找到对应 .o 文件路径
加载 .o（首次加载；后续从缓存取）
组装 kernel 参数：
- TilingData（tiling 阶段产出）
- self 的 NPU 地址
- result 的 NPU 地址
- workspace 地址
下发 kernel launch 指令到 NPU

第五层：AscendC Kernel（NPU 上执行）

AI Core × N（并行）：

CopyIn：从 Global Memory 搬数据到 UB（Local Memory）
Compute：向量计算单元执行 Relu（max(x, 0)）
CopyOut：结果从 UB 搬回 Global Memory

所有核完成 → NPU 通知 CPU → result tensor 可用。

图模式链路

PyTorch (用户代码)
  └─ torch_npu / op-plugin（同上，图模式入口）
       └─ torch_air（PyTorch 图捕获与 Ascend 编译框架）
            └─ GE（Graph Engine：图优化、算子融合、子图划分）
                 └─ CANN Runtime（逐算子调度，可复用 aclnn 路径）
                      └─ AscendC Kernel binary

图模式完整执行链路分两个阶段：编译期（一次） 和 执行期（每次推理）。

编译期

用户代码：

model = torch.compile(my_model)   # 或 torch_npu 的图模式接口

第一层：PyTorch 图捕获

torch.compile 触发 FX Graph 追踪
捕获所有算子及其依赖关系
产出 PyTorch FX IR（计算图）

第二层：torch_air

接收 FX IR
将 PyTorch 图转换为 Ascend GE IR
传入 GE

第三层：GE（Graph Engine）编译

GE 对图执行：

图优化
- 算子融合（如 Conv+BN+Relu → 单个 kernel）
- 常量折叠、冗余消除
InferShape
- 调用 libophost_ops_nn.so 中的 InferShape 函数
- 为图中每个算子推导输出 shape
- 确定每个 tensor 的大小，用于显存规划
显存规划
- 静态分配图中所有 tensor 的 NPU 显存
Kernel 选择
- 每个算子查 config/*.json，确定使用哪个 .o
Tiling
- 静态 shape 时在此预计算
- 调用 libophost_ops_nn.so 中的 tiling 函数
- TilingData 固化进执行计划

最终产出：编译好的 GE Graph（序列化执行计划，可反复执行）。

执行期

用户代码：

out = model(input)   # 第 N 次推理，直接复用编译结果

第一层：GE Runtime

加载编译好的 GE Graph
动态 shape 时重新调用 tiling 函数，更新 TilingData
静态 shape 时直接使用编译期预计算的 TilingData

第二层：按图拓扑顺序下发 kernel

对图中每个算子（或融合后的算子）组装参数：
- TilingData
- 输入地址
- 输出地址
- workspace 地址
下发 kernel launch 指令到 NPU

第三层：NPU 执行

AI Core × N（并行）执行：CopyIn → Compute → CopyOut
图中多个 kernel 可流水线重叠执行
所有算子完成后，输出 tensor 可用

编译过程

编译 kernel binary

bash build.sh --opkernel --ops=index_put_with_sort_v2 --soc=ascend950

阶段一：CMake Configure

源码目录	内容	归属 target
`op_host/`	tiling、infershape	ophost
`op_host/arch35/`	该 SoC 架构的 tiling 变体	ophost（SoC 分支）
`op_api/`	`_def.cpp`（算子接口定义）	opbuild / opapi
`op_kernel/`	AscendC kernel 源码	binary 编译输入

仅生成 Makefile，不编译。

阶段二：cmake --build，依次执行以下 target：

Target 1 — ophost：编译 tiling / infershape 动态库

g++ -std=c++17 ... op_host/**/*.cpp
    → build/libophost_ops_nn.so

Target 2 — opbuild_gen_aclnnExc：从 `_def.cpp` 生成算子描述 `.ini`

op_build 处理 op_api/index_put_with_sort_v2_def.cpp
    → build/tbe/config/ascend950/index_put_with_sort_v2.ini

Target 3 — ascendc_impl_gen：生成 `asc_opc` 的输入 Python 描述文件

读取 .ini + op_kernel/*.cpp
    → build/tbe/index_put_with_sort_v2_xxx.py

Target 4 — gen_bin_scripts：按 dtype 变体生成编译脚本

读取 op_host/config/ascend950/index_put_with_sort_v2_binary.json

asc_opc build/tbe/index_put_with_sort_v2.py \
  --main_func=index_put_with_sort_v2 \
  --input_param=.../float16_params \
  --soc_version=Ascend950A \
  --output=build/binary/ascend950/index_put_with_sort_v2/ \
  --impl_mode=high_performance

Target 5 — binary：多线程并行编译所有 dtype 变体

asc_opc（底层调用 BiSheng 编译器）：

输入：.py 描述 + dtype 参数
产物：
- IndexPutWithSortV2_float16.o
- IndexPutWithSortV2_float16.json
- IndexPutWithSortV2_float32.o
- IndexPutWithSortV2_float32.json

编译 aclnn 接口层

bash build.sh --opapi --soc=Ascend950

Target — opapi_nn：编译 aclnn 接口动态库（纯 CPU 侧，编译时与 SoC 无关）

opbuild 从 _def.cpp 自动生成 aclnn_index_put_with_sort_v2.cpp
    + 所有其他算子的 aclnn_xxx.cpp
        ↓ g++ -std=c++17
    → build/libopapi_nn.so

假设要实现一个 `relu` 函数

第一层：op_plugin

这是离 PyTorch 最近的一层，需要写一个 wrapper，做两件事：

向 PyTorch Dispatcher 注册 NPU 实现
调用 aclnn

op_plugin/ops/opapi/ReluKernelNpuOpApi.cpp

#include "op_plugin/OpApiInterface.h"
#include "op_plugin/utils/op_api_common.h"
 
// 1. wrapper 实现
at::Tensor relu_npu(const at::Tensor& self) {
    // InferShape + 分配输出 tensor
    at::Tensor result = npu_preparation::apply_tensor(self);
 
    // 调用 aclnn（宏内部处理 GetWorkspaceSize + 执行）
    EXEC_NPU_CMD(aclnnRelu, self, result);
 
    return result;
}
 
// 2. 注册到 PyTorch Dispatcher
TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {
    m.impl("relu", relu_npu);
}

第二层：torch_npu

不需要改动。 torch_npu 提供基础设施（npu_preparation、stream 管理、device 管理等），op_plugin 直接使用这些工具，不需要在 torch_npu 里注册或添加任何东西。

第三层：CANN 算子（前面说的 6 个文件）在 ops 仓：

op_api/relu_def.cpp                            ← OpDef，生成 aclnn 接口
op_host/relu_tiling.h                          ← TilingData 结构体
op_host/relu_tiling.cpp                        ← Tiling 函数
op_host/config/ascend950/relu_binary.json      ← 编译哪些 dtype
op_kernel/relu.h                               ← Kernel 类
op_kernel/relu.cpp                             ← __global__ 入口

fxlin's Blog

目录

编译执行链路

执行过程

单算子

图模式链路

编译期

执行期

编译过程

编译 kernel binary

Target 1 — ophost：编译 tiling / infershape 动态库

Target 2 — opbuild_gen_aclnnExc：从 _def.cpp 生成算子描述 .ini

Target 3 — ascendc_impl_gen：生成 asc_opc 的输入 Python 描述文件

Target 4 — gen_bin_scripts：按 dtype 变体生成编译脚本

Target 5 — binary：多线程并行编译所有 dtype 变体

编译 aclnn 接口层

假设要实现一个 relu 函数

目录

Target 2 — opbuild_gen_aclnnExc：从 `_def.cpp` 生成算子描述 `.ini`

Target 3 — ascendc_impl_gen：生成 `asc_opc` 的输入 Python 描述文件

假设要实现一个 `relu` 函数