当前位置：首页 > news >正文

TBE 算子开发框架解析

news 2026/6/10 15:55:41

前言

2019年刚开始做昇腾算子开发，用的是TBE框架，用DSL写算子，框架自动生成调度策略和代码。当时觉得"自动生成"很高大上，后来出了Ascend C，才发现TBE有性能天花板——自动生成的调度策略不如手写优化。

很多人以为TBE过时了，其实它还在广泛用于算子的快速开发。TBE开发效率比Ascend C高3倍，性能差10-15%，适合对性能要求不极致的场景。

TBE 的定位

TBE（Tensor Boost Engine）是CANN早期的算子开发框架，用DSL（Domain Specific Language）描述算子计算，自动生成调度策略和底层代码。

CANN 架构中的 TBE： AscendCL（编程接口层） ↓ AOL 算子库（ops-math/ops-nn/...） ↓ GE（图引擎） ↓ TBE 算子开发框架 ← 你在这（已逐步被 Ascend C 替代） ↓ 生成 调度策略 + 底层代码（C++/CUDA） ↓ Runtime（运行时） ↓ 驱动层

TBE 已被 Ascend C 逐步替代，但还在广泛用于已有算子的维护和快速开发。

工程经验：不复用TBE用Ascend C重写算子，开发周期多2-3周，性能提升10-15%。如果算子已经用TBE写好了，没必要重写，除非性能是刚需。

TBE 的核心技术

1. DSL 描述计算

TBE用DSL描述算子计算，不需要手写调度策略。

# TBE DSL：写一个 MatMul 算子fromtbeimportdsl@dsl.register_op("MatMul")defmatmul(A,B,C,M,K,N):# 描述计算（不需要写调度）foriinrange(M):forjinrange(N):C[i,j]=0forkinrange(K):C[i,j]+=A[i,k]*B[k,j]returnC# 自动生成调度策略 + 底层代码dsl.auto_schedule(matmul)dsl.build(matmul,target="npu")

对比 Ascend C（手写调度）：

// Ascend C：手写 MatMul 算子（200行）#include"kernel_operator.h"classMatMulKernel{public:__aicore__voidProcess(GM_ADDR a,GM_ADDR b,GM_ADDR c,intM,intK,intN){// 手写 TilingconstexprintTILE_M=64,TILE_K=64,TILE_N=64;// 手写缓存管理TPipe pipe;TBuf<TPosition::A1>A_L0A;TBuf<TPosition::B1>B_L0B;TBuf<TPosition::C1>C_L0C;pipe.AllocBuf(A_L0A,TILE_M*TILE_K*sizeof(half));pipe.AllocBuf(B_L0B,TILE_K*TILE_N*sizeof(half));pipe.AllocBuf(C_L0C,TILE_M*TILE_N*sizeof(half));// 手写流水线for(intm=0;m<M;m+=TILE_M){for(intn=0;n<N;n+=TILE_N){InitC(C_L0C,TILE_M,TILE_N);for(intk=0;k<K;k+=TILE_K){// 手写搬运DataCopy(A_L0A,a+m*K+k,TILE_M*TILE_K*sizeof(half));DataCopy(B_L0B,b+k*N+n,TILE_K*TILE_N*sizeof(half));// 手写矩阵乘MatMul(C_L0C,A_L0A,B_L0B,TILE_M,TILE_K,TILE_N,{.accumulate=true});}// 手写写回DataCopy(c+m*N+n,C_L0C,TILE_M*TILE_N*sizeof(half));}}}};

DSL 20行 vs Ascend C 200行，开发效率高10倍。

2. 自动调度策略

TBE自动分析DSL代码，生成最优调度策略（Tiling、缓存管理、流水线）。

自动Tiling：

# TBE 自动 Tiling（根据 L0A/L0B/L0C/L1 容量）fromtbeimportdsl@dsl.register_op("MatMul")defmatmul(A,B,C,M,K,N):# ... 计算描述# 自动 Tiling（不需要手写）# TBE 自动算最优 tile_m/tile_k/tile_n# 考虑 L0A/L0B/L0C/L1 容量约束# 考虑 MAC 阵列利用率pass# 查看自动生成的 Tiling 参数sch=dsl.auto_schedule(matmul)print(sch.tiling)# 输出：# tile_m = 64# tile_k = 64# tile_n = 64

自动缓存管理：

# TBE 自动缓存管理（不需要手写 AllocBuf）sch=dsl.auto_schedule(matmul)# 查看自动生成的缓存管理代码print(sch.cache_manager)# 输出：# TPipe pipe;# TBuf<TPosition::A1> A_L0A = pipe.AllocBuf(64*64*2);# TBuf<TPosition::B1> B_L0B = pipe.AllocBuf(64*64*2);# TBuf<TPosition::C1> C_L0C = pipe.AllocBuf(64*64*2);

自动流水线：

# TBE 自动流水线（不需要手写 double buffer）sch=dsl.auto_schedule(matmul)# 查看自动生成的流水线代码print(sch.pipeline)# 输出：# // Double Buffer# half A_L0A_0[64][64], A_L0A_1[64][64];# half B_L0B_0[64][64], B_L0B_1[64][64];# // Pipeline: Cube 算当前 tile，DMA 搬下一个 tile

3. 自动代码生成

TBE自动生成底层C++/CUDA代码，可以直接编译运行。

# 生成底层代码dsl.build(matmul,target="npu",output="matmul_kernel.cpp")# 查看生成的代码catmatmul_kernel.cpp# 输出（自动生成的 C++ 代码）：# #include "kernel_operator.h"## class MatMulKernel {# public:# __aicore__ void Process(GM_ADDR a, GM_ADDR b, GM_ADDR c,# int M, int K, int N) {# // 自动生成的 Tiling# constexpr int TILE_M = 64, TILE_K = 64, TILE_N = 64;## // 自动生成的缓存管理# TPipe pipe;# TBuf<TPosition::A1> A_L0A;# ...# }# };

工程经验：TBE自动生成的代码，性能比手写Ascend C差10-15%。原因是自动调度策略是"通用最优"，不是"特定算子最优"。比如MatMul自动生成的Tiling是tile_m=64，但特定shape（M=1）最优是tile_m=1（用Vector Unit算）。

TBE vs Ascend C 对比

维度	TBE（DSL）	Ascend C（C++）
开发效率	高（20行DSL）	低（200行C++）
性能	85-90%	100%
学习曲线	缓（只懂Python即可）	陡（要懂C++、ACL、内存管理）
调试	易（Python堆栈清晰）	难（Core Dump难定位）
适用场景	快速开发、维护已有算子	性能极致优化、新算子开发

工程经验：新算子开发用Ascend C（性能极致），维护已有算子用TBE（开发效率高）。不要混用（TBE生成的代码再手写优化，调试很难）。

TBE 的使用流程

1. 安装 TBE

# TBE 已内置在 CANN 里，不需要单独安装# 确认 TBE 可用python-c"from tbe import dsl; print('TBE available')"# 输出：# TBE available

2. 写 DSL 算子

# matmul_dsl.pyfromtbeimportdsl@dsl.register_op("MatMul")defmatmul(A,B,C,M,K,N):# 描述计算foriinrange(M):forjinrange(N):C[i,j]=0forkinrange(K):C[i,j]+=A[i,k]*B[k,j]returnC# 自动调度 + 代码生成sch=dsl.auto_schedule(matmul)dsl.build(matmul,target="npu",output="matmul_kernel.cpp")

3. 编译运行

# 编译生成的 C++ 代码npu-smiset-tmm-s0-dmatmul_kernel.o matmul_kernel.cpp# 链接成动态库ld-sharedmatmul_kernel.o-olibmatmul.so# 在 ACL 中调用aclError ret=aclrtLaunchKernel(matmul_kernel, grid, block, args,0, stream);