8 月 3 日，2026 夏季 InfiniTensor 训练营 CUDA 编程方向第七课《寒武纪软件栈 NeuWare 和 BANG 异构计算平台 2》开讲。
上节课《寒武纪软件栈 NeuWare 和 BANG 异构计算平台 1》讲解了寒武纪 NeuWare 软件栈架构与 BANG 异构计算平台，涵盖 PyTorch/TensorFlow 框架适配、CNToolkit 工具链、算子库开发及性能调优等核心知识点。

本节课聚焦 BANG 异构计算平台，通过 Sigmoid 和 MatMul 两个案例，串联了 MLU 编程的核心概念。带大家掌握 PyTorch 自定义算子开发流程及性能调优思路。

算子开发与集成实验

实现 Sigmoid 算子

1. 工程目录结构解析

• 顶层：扩展算子代码目录 + 构建脚本 + 测试用例
• 源码分层：
  • mlu/：存放 BANG C++ 实现的算子 Kernel 及 API
  • cpp/：PyTorch C++ API 封装层
  • .py：Python 接口包装层

2. BANG C++ Kernel 开发核心

1.Built-in 接口调用：使用 CNCC/BANG C++ 手册中的 active_sigmoid 向量 SIMD 接口

2.Kernel定义：使用 __mlu_global__ 修饰符，模板化支持多数据类型

3.任务拆分策略：

• 利用内置变量 taskDim 进行总元素拆分（虚拟 Task 映射到物理核/Cluster）
• 多级拆分：先按多核多任务拆分，再在单任务内用循环处理
• 边界处理：针对非对齐数据，增加尾部余数段处理代码

4.异步访存模型（高性能关键）：

• Async 流程：async copy (DRAM→NRAM) → sync IO → Compute → sync Compute → async copy (NRAM→DRAM) → sync IO
• 入门建议：初期可省略 async/sync，依赖 CNCC 编译器自动管理依赖；功能验证无误后，再切换至显式异步控制以追求极致性能

集成进 PyTorch

1.C++封装：将 Kernel 指针/大小参数转换为 PyTorch Tensor 对象

2.算子注册：使用 PyTorch 官方宏将自定义算子注册到分发机制中

3.混合编译流程：

算子性能优化实验

原始标量版本矩阵乘

1. 直接将 CPU 标量代码移植为 MLU Kernel
2. 单核单 Block 执行（1×1×1配置）
3. 使用 notifier 挂在队列进行精准耗时统计
4. 结果：性能极差，仅作为 Baseline

使用片上 NRAM 的矩阵乘

1. 优化原理：利用片上 NRAM（几百K~2MB）的数据局部性，降低访存延迟
2. 编程约束：NRAM 属于片上高速存储空间，需手动申请静态数组并按块拆分
3. 访存模式：两读一写，数据搬运至 NRAM 后进行标量计算
4. 结果：2000ms+ → 96ms（仍有差距）

向量化计算替换

1.将标量循环替换为 BANG Built-in 函数 __bang_matmul

2.释放 MLU 张量/向量运算单元能力

使用多核的矩阵乘

1.Host端：增大 Block 数量（如 128），匹配数据规模

2.Device端：

• 总任务按 Block 数切分 + 循环迭代
• 循环不变量外提：W 矩阵复用，将其从 Global DRAM 到 NRAM 的搬运移至循环外（编译器亦可自动优化）
• 循环内仅切片访问左矩阵并计算写回

3.结果：对比 46008ms，MLU 仅需0.595ms

使用三级流水的矩阵乘

1.依赖控制：

• 显式 Async 写法：区分 sync_io / sync_compute / sync_move
• 简化写法：使用全 sync 保证正确性（牺牲少量性能换取开发效率）

2.实战备注：

• 本例因数据规整且编译器已做自动优化，显式流水收益有限
• 真实场景：面对可变 Shape、NRAM 容量限制等复杂情况，手动流水与精细化内存管理是必备技能

加入训练营

📺 完整课程内容，请观看直播或查看课程回放：

InfiniTensor 官网： https://www.infinitensor.com

B 站直播： InfiniTensor 官方直播间

视频号直播： InfiniTensor 视频号预约

答疑交流： 训练营官方社群 ⬇️

关注与交流