Keyboard shortcuts

Press or to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

English | 中文版

5. 规模化:覆盖 MultiKernelBench 全部类别的 508 个内核

在单一基准测试和等价验证之外,我们系统性地扩展了 ascend-rs 的内核覆盖范围,实现了对 MultiKernelBench 基准套件全部 300 个 PyTorch 参考内核的完整 1:1 覆盖,涵盖 17 个类别(激活函数、网络架构、注意力机制、广播运算、卷积、融合算子、索引操作、损失函数、数学运算、矩阵乘法、归一化、优化器、池化、归约、缩放、分块、多核)。

ascend-rs 目前包含 505 个 Rust NPU 内核,全部可通过 MLIR 代码生成后端编译。这些内核按验证层级分为以下级别:

  • 16 个可部署内核 — 通过完整的 Rust→MLIR→C++→bisheng 流水线编译,已部署到 NPU 硬件上执行
  • 413 个测试在 Ascend 910B3 上通过 NPU 正确性验证 — 在真实硬件上与 CPU 参考验证,0 失败、0 崩溃;代表性内核(第 4.5 节)与手写 AscendC C++ 逐位相同。包含 37 个矩阵乘法测试通过 CANN 的 aclnn 算子 API(aclnnMm、aclnnAdd、aclnnAddmm、aclnnRelu、aclnnMul、aclnnReduceSum)执行,以及全部卷积、池化、缩放、索引和优化器内核
  • 486 个编译测试内核 — 已验证可通过 MLIR 后端编译并通过 CPU 级正确性测试

Cube 引擎矩阵乘法内核——此前因混合 AIV/AIC 二进制中 TPipe L1/CBUF 队列分配问题而受阻——现已通过 CANN 内置算子 API 正确执行。两阶段 aclnn 算子模式(GetWorkspaceSize + Execute)从 libopapi.so 动态加载,完全绕过自定义内核编译,利用 Cube 引擎的内置优化算子。组合算子链(如 aclnnMm + aclnnRelu + aclnnAdd 实现 ResNet 残差块)使融合矩阵乘法变体得以实现,否则需要自定义 Cube 内核开发。

类别内核数实现方式
激活函数 (16)relu、sigmoid、gelu、tanh、softmax、elu、selu、swish、mish、softplus、softsign、hardsigmoid、hardswish、leaky_relu、log_softmax、gelu_tanh向量指令 + kernel_ops 组合算子
网络架构 (41)AlexNet/VGG/ResNet 全连接层、DenseNet 块、MobileNet/EfficientNet、ViT/Swin MLP、MinGPT、LSTM 门控/单元、GRU 门控、Mamba SSM矩阵乘法 + 激活 + 归一化组合
注意力机制 (15)缩放点积、因果、交叉、多查询、分组查询、KV 缓存、跨模态、线性、稀疏、窗口因果、SwiGLU、GeGLU、掩码填充缩放 + 掩码 + softmax 模式
广播运算 (8)add_bias、逐元素乘/除/减/最大/最小、clamp、平方二元向量指令
卷积 (34)标准 conv2d、深度可分离 conv2d、转置 conv2d 变体标量嵌套循环(不使用 Cube 引擎)
融合算子 (86)matmul+gelu、gemm+relu+divide、norm+激活、多算子链(3-6 个算子融合)链式向量指令 + 流水线屏障
索引操作 (12)gather、scatter、scatter_add、index_select、index_copy、index_add、embedding、masked_fill、inplace_update、take_along_dim标量嵌套循环 + 边界检查索引
损失函数 (6)MSE、Huber、hinge、余弦相似度、交叉熵、KL 散度归约 + 算术
数学运算 (5)累积和(3 种变体)、累积积、矩阵标量乘法标量循环 + 向量操作
矩阵乘法 (17)标准、批量、对称、带偏置、缩放、GEMM、宽矩阵、累加、对角缩放、外积Cube 引擎(Mmad FFI)
归一化 (9)layernorm、rmsnorm、batch/group/instance norm、L1/L2/Frobenius 范数归约 + 归一化模式
优化器 (6)SGD、SGD+动量、Adagrad、RMSprop、Adam、扩展变体原地缓冲区算术
池化 (6)全局平均/最大/最小池化、融合池化+sigmoid、LP 池化基于归约
归约 (5)最大、最小、求和、均值、乘积硬件归约指令
缩放 (5)最近邻、线性插值、双三次权重、加权求和、三线性插值算术
分块 (16)256 元素分块的激活函数和运算变体循环 + 分块缓冲区分配
多核 (16)AICore 块级并行变体get_block_idx() 工作分配

为支持这一广度,我们在 kernel_ops.rs 中新增了 17 个组合算子——如 elu_f32mish_f32rms_norm_f32mse_loss_f32cosine_similarity_f32——每个都由基础向量指令组合而成,并正确放置流水线屏障。

卷积和索引/gather/scatter 类别通过标量嵌套循环模式实现,在 API 层面达成 MultiKernelBench 的完整覆盖。CPU 正确性测试(cargo test -p kernel_correctness)验证了涵盖所有类别的 80 个代表性内核的数值精度。其余编译测试验证了通过 MLIR 后端的成功编译,但未进行 CPU 级数值检查。

进度报告 — 截至当前代码库的验证状态(通过 count_kernels.sh 和硬件测试日志确认):

验证层级数量说明
编译测试通过486通过 MLIR 后端编译 + CPU 级正确性(cargo test -p compiletest
910B3 正确性验证413在 Ascend 910B3 上通过 NPU 正确性测试(0 失败、0 崩溃);包含 37 个矩阵乘法(aclnn)、全部卷积/池化/缩放/索引/优化器内核
与 AscendC 性能对等4开销 ≤2%(第 4.3–4.4 节):softmax、relu、sigmoid、tanh
可部署(完整流水线)16通过 Rust→MLIR→C++→bisheng 编译并在 NPU 上执行
内核总数505全部可通过 MLIR 代码生成后端编译

413 个通过 NPU 正确性测试的测试覆盖所有内核类别:向量指令内核(激活函数、归约、融合算子链、多核并行)、Cube 引擎矩阵乘法(通过 aclnn 算子组合)、卷积、池化、缩放、索引操作和优化器——0 失败、0 崩溃。