Assignment 4(昇腾版): Programming a Machine Learning Accelerator

本文是 README.md(Trainium2 版)的昇腾对照版:章节顺序、小节划分、知识点一一对应,只把平台换成华为昇腾。

  • 目标平台:华为昇腾 CANNLab 的一个 NPU,机型为 A2 = Ascend 910B 或 A3 = Ascend 910C

  • 编程语言:Ascend C(C/C++ 扩展),不是 NKI

  • 性能口径:Trainium2 版给的所有阈值与数字(128 lanes、<700 μs、4596/1002 μs)在昇腾上无意义,必须自建基线。那些数字是针对 Trainium2 硬件实测的绝对值(性能用例 fp16 满分线 1002 μs 反推约 58 TFLOPS,是它 79 TFLOPS 峰值的 73%),而昇腾的芯片、片上存储(L0C 128 KB vs PSUM 2 MiB)与工具链都不同,照搬没有意义。

    本手册的任务、评分点与判定口径是完整的;只有"多少微秒算达标"这一类绝对数字需要重新标定,方法见 Part 2 的 Testing。


Overview

在本作业中,你将学习如何为华为昇腾(Ascend)AI 处理器实现并优化 kernel。与 Trainium2 一样,昇腾也提供多个面向张量/向量的加速计算引擎,以及由软件管理的片上存储,为这些引擎提供高带宽的数据通路。

作业同样分为两部分。Part 1 通过向量加的几个实现熟悉昇腾的架构与数据搬运模式,并自己写一个矩阵转置 kernel;Part 2 在昇腾上实现一个融合的 convolution + maxpool 层。

总体而言,本作业会让你:

1) 体验张量处理与片上 SRAM 管理的底层细节。

2) 理解循环分块(loop blocking)与循环融合(loop fusion)这类保局部性优化的价值。

先说明一处结构性差异:Trainium2 的 Part 1 在单个 NeuronCore 上编程,"多核"只是架构背景板;昇腾是 SPMD 多核模型,一段 kernel 会被很多个 AI Core 同时执行,blockDim 决定并行度。因此在昇腾上,"每一级分块"的问题里天然多了一层:核间怎么切。

Environment Setup

你将在华为昇腾云开发环境(CANNLab)上编程与测试你的代码。环境申请与连接的要点是:

  • 入口为 gitcode.com/org/cann/cannlab,外部用户有免费 100 卡时,关机不计时;
  • 本作业是单算子调优,务必选单卡/单 NPU 规格(选 8 卡会按卡的倍数扣卡时);
  • 连接方式为 WebIDE 或本地 VS Code。

登录之后,用课程 Github 下载作业起始代码:

git clone https://github.com/stanford-cs149/asst4-trainium2

进入 asst4-trainium2 目录后,Trainium2 版让你跑 source install.sh,昇腾侧不需要跑它。原因是这个脚本只做三件事,而在昇腾云环境里这三件事都已经被平台或工具链替代:

install.sh 做的事 昇腾侧的对应物
激活带依赖的 Python venv,并写进 ~/.bashrc CANN 工具链预置在镜像里,登录即在 PATH 中;用 source $ASCEND_HOME_PATH/bin/setenv.bash 显式加载环境变量
安装并启动 InfluxDB、influx setup 建库 没有对应物:昇腾的剖析工具 msprof 与 MindStudio Insight 直接产出本地文件,不需要数据库、不需要转发端口
(隐含)SSH 转发 3001 / 8086 才能看 neuron-profile GUI 不需要:msprof 的结果落盘为本地目录,用 MindStudio Insight 或 msprof 自带的统计输出查看

取而代之,环境准备只需要确认下面三件事:

npu-smi info                      # 看卡与 die(NPU ID 会随会话变化,别硬编码 -i)
echo $ASCEND_HOME_PATH            # CANN 根目录
cat $ASCEND_HOME_PATH/version.info # CANN 版本

[!NOTE] npu-smi 的输出里,Chip Name + NPU Name 拼起来才是 --soc-version(例如 Ascend910_9362)。A2 单 die、A3 双 die,本作业一次 kernel 只用到一个 NPU 的一个 die。

Part 0: Getting familiar with Ascend and AI Core Architecture

Ascend Architecture Overview

先认识昇腾。

本作业在 CANNLab 上用的是 A2 = Ascend 910B 或 A3 = Ascend 910C 的一个 NPU。一颗 die 上有很多个 AI Core,它是昇腾的算力单位,地位等同于 NVIDIA 的 SM,也对应 Trainium 的 NeuronCore(注意:对应的是 NeuronCore,不是整颗芯片)。本作业的 kernel 就跑在 AI Core 上,blockDim 决定同时用几个核。

NeuronCore-v3 把 Tensor/Vector/Scalar 三种引擎放在同一个核里。A2/A3 采用"分离模式":矩阵计算单元部署在 AIC(Cube Core),矢量计算单元部署在 AIV(Vector Core)。两边各有自己的 Scalar 单元、独立取指、独立代码段,再按 1:N 的比例组合成一个 AI Core。本作业实测 A3 的 Ascend910_9362 为 AIC = 20,AIV = 40,即 1 : 2。这个数字请用 GetCoreNumAic() / GetCoreNumAiv() 运行时查询,不要硬编码。

AI Core硬件架构的抽象

图 1 AI Core 的分离模式(A2 / A3)

一个 AI Core 由三类部件构成:

组件分类 组件 职责
计算单元 Cube(AIC) 矩阵运算。一条指令完成两个 16×16 fp16 矩阵乘(分形 16×16×16)
计算单元 Vector(AIV) 向量运算。SIMD,源与目的都必须在 UB 中,要求 32 B 对齐
计算单元 Scalar 可看作核内小 CPU:循环、分支、地址计算,并把向量/矩阵/搬运指令发射到对应队列
存储单元 Local Memory 片上存储总称,对应编程模型里的 LocalTensor:UB / L1 / L0A / L0B / L0C / BT / FP
搬运单元 DMA Global Memory 与 Local Memory 之间、以及各级 Local Memory 之间的数据搬运

更细的硬件说明(计算单元、存储单元、搬运单元的功能与结构)见官方《Ascend C 基本架构》。

Atlas A2 训练系列产品硬件架构图

和 NeuronCore 最容易被忽略的一处差异:NeuronCore-v3 的三种引擎在同一核内直接协作;昇腾把 Cube 和 Vector 拆成两个核,AIC 与 AIV 不能直接互访对方的片上 Buffer(L1/L0C 对 AIV 不可见,UB 对 AIC 不可见)。

Ascend C 抽象硬件架构

上面那张是硬件视角。写 kernel 时用的抽象只有三个概念,正是这层抽象把 A2/A3 的差异屏蔽掉了,所以同一份 kernel 代码换机型不用改(只需改并行度):

  • Global Memory(GM) → GlobalTensor<T>:显存里的数据,GM_ADDR 入参;
  • Local Memory → LocalTensor<T>:片上 Buffer,用 TQue / TBuf 申请,TPosition 指定放在哪一级;
  • DMA → DataCopy 系列接口:搬入搬出。

LocalTensor 是UB 上的一段逻辑数据视图,承载数据类型 T(如 half、float)和长度信息,本身不直接分配内存,需从 TBuf 或 TQue 中获取。

TPipe          ← 统一管理者,管 UB 内存分配 + 同步事件
  ├── TQue     ← 带队列语义的缓冲区(有同步)
  └── TBuf     ← 纯静态缓冲区(无同步)
        ↓ 两者都产出
   LocalTensor<T>   ← 你真正拿来算的东西

类比一下:TPipe 是物业,TQue/TBuf 是两套不同租约的房间,LocalTensor 是你住进去之后拿到的钥匙。

图 2 Ascend C 抽象硬件架构:三条流

理解这张图的关键是三条流:

  1. 指令流(蓝):Scalar 读取指令序列,把向量计算、矩阵计算、数据搬运指令分别发射到各自独立的指令队列,各队列异步并行执行。这让 MTE2 搬下一块数据、Cube/Vector 算当前块、MTE3 搬出上一块可以同时进行,也就是流水线并行。
  2. 同步信号流(绿):队列之间仍有数据依赖(比如"搬完才能算"),Scalar 会下发同步指令。编程模型(TQue 的 EnQue/DeQue)已替你插入同步;只有手写 PipeBarrier / SetFlag-WaitFlag 时才需要自己管,官方也推荐用编程范式而不是手写同步。
  3. 数据流(橙):有两条典型路径:

    • Vector 路径:GM → UB → Vector → UB → GM
    • Cube 路径:GM → L1 → L0A/L0B → Cube → L0C → Fixpipe → GM / L1

      注意 Cube 路径比 NKI 多出 L1 → L0A/L0B 这一级显式 staging,这是 Part 2 性能开销的主要来源之一。

搬运单元按通路分工,记这张表就够了:

搬运单元 负责通路 备注
MTE1 L1 → L0A / L0B、L1 → BT Buffer 片上到片上,按分形搬运
MTE2 GM → {L1, L0A/L0B}、GM → UB 搬入;按分形 / 按 Cache Line 更优
MTE3 UB → GM 搬出
Fixpipe L0C → {GM, L1}、L1 → FP Buffer 搬出,可随路做量化/反量化、Relu、NZ→ND 格式转换

MTE = Memory Transfer Engine。「内存传输引擎」,也常被叫作 LSU(Load Store Unit)。

(MTE1/2/3 与 Fixpipe 的可用性、以及"随路转换"的具体能力随架构版本不同,A2/A3 属于 NPU 架构版本 220x。)

流水线、TQue 队列与双缓冲(BUFFER_NUM)

昇腾的片上内存是显式管理的:数据必须由指令搬进 UB 才能算。一个 AI Core 上有三种"引擎"——MTE2(GM → UB 搬入)、Vector / Cube(计算)、MTE3(UB → GM 搬出);它们之间靠队列(TQue / TPipe)通信与同步:

接口 角色
AllocTensor / FreeTensor 借 / 还 buffer(资源占有与释放)
EnQue / DeQue 生产 / 消费(谁准备好了、谁等谁)
(底层)SetFlag / WaitFlag 硬件同步指令:MTE2 写完才放行 Vector

队列深度 ≠ 双缓冲(TQue 简介 :"这里的队列深度和 double buffer 无关,队列机制用于实现流水线并行,double buffer 在此基础上进一步提高流水线的利用率")。它们控制流水线的两个不同属性:

设置在哪里 控制什么 官方建议
TQue<TPosition::VECIN, depth> 的模板参数 depth 同步窗口:连续 EnQue 几次才算一次完整生产 恒为 1(编译器对 depth=1 有专门优化;仅 preload 场景用 2,"其他情况均不推荐 ≥ 2")
pipe.InitBuffer(que, num, len) 的 num 每队列的内存块份数:num = 2 就是双缓冲 双缓冲推荐 2;"即使队列的深度为 1,仍可以开启 double buffer"

双缓冲在做什么:让 MTE2 提前搬下一份、Vector 算当前份、MTE3 搬出上一份,三台引擎重叠起来:

MTE2:   [x0] [x1] [x2] ...
Vector:      [z0] [z1] ...
MTE3:             [z0] ...

本项目代码里用BUFFER_NUM变量,控制的就是 InitBuffer 的 num:流水线上同时在途的 tile 份数(每份 buffer 占一个同步事件槽)。加深 num(3/4/…)=加深预取深度、吸收 GM 延迟抖动——但没有"第 4 台引擎",并行度上限是 3 台引擎,收益边际递减,还要额外占 UB。

上限:每个 buffer 存一个同步事件 eventID,硬件按存储位置(TPosition)分配事件池;A2 / A3 每个 TPosition 事件数为 8,所以同一 TPosition 上所有队列的 buffer 总数 ≤ 8(队列多了各自都要少)。加上 UB 容量因素,本作业 --kernel stream 用 num = 2 正是官方推荐配置;且本项目 TQue 的 depth 固定为 1(TQue<TPosition::VECIN, 1>),双缓冲只由 InitBuffer 的 num 控制。

TPosition(逻辑存储位置):TQue 的第一个模板参数告诉 TPipe 从哪块片上存储分配内存(TPosition):

TPosition 物理存储 用途
VECIN / VECOUT / VECCALC UB 矢量计算的搬入 / 搬出 / 临时
A1 / B1 L1 Cube 的整块 A / B 矩阵
A2 / B2 L0_A / L0_B Cube 切块后的 A / B
CO1 / CO2 L0_C(Cube Out) Cube 的结果 C
GM HBM 外部存储

同步事件池按 TPosition 划分:同一位置的队列共用一池(A2/A3 各 8),这就是"同一个 TPosition 上 TQue Buffer 总数有上限"的含义。

Ascend Memory Hierarchy

在 Assignment 3 里,GPU 的存储层次是 host 内存 → device global memory → 每 block 的 shared memory → 每线程私有寄存器。

昇腾的层次结构是:片外 GM(HBM)→(L2 Cache)→ 片上 Local Memory,其中片上又分成两条互不相通的路径:

  • AIV 路径:Unified Buffer (UB) 192 KB,Vector / Scalar 的输入输出;
  • AIC 路径:L1 Buffer 512 KB(Cube 的中转区)→ L0A / L0B 各 64 KB(左/右矩阵输入)→ L0C 128 KB(矩阵乘输出兼累加器);另有 BT Buffer 1 KB(bias)与 FP Buffer 2 KB(量化/激活参数)。

图 3 Ascend 存储层次与数据流(容量为架构级常数)

层级 容量 角色 对齐 / 格式要求
GM(HBM) A3 单 die 64 GB 片外主存,所有 AI Core 共享 —
L2 Cache 168 MiB(A2 为 96 MiB) 硬件自动缓存所有 GM 访问,编程时不可见 Cache Line 128/256/512 B
L1 Buffer 512 KB Cube 的数据中转区,暂存反复使用的数据 32 B;推荐 FRACTAL_NZ
L0A / L0B 各 64 KB Cube 左矩阵 / 右矩阵输入 512 B;分别 FRACTAL_ZZ / FRACTAL_ZN
L0C 128 KB Cube 输出兼累加器 64 B;FRACTAL_NZ
Unified Buffer 192 KB Vector / Scalar 的输入输出 32 B;对格式无要求
BT / FP Buffer 1 KB / 2 KB BiasTable 存放矩阵计算中的Bias。

/ Fixpipe 参数: 放量化参数、Relu参数等。
64 B

与 Trainium2 的量级对照(这是本作业最需要提前认清的一条):

NeuronCore-v3(Trainium2) 昇腾 A2/A3 差距
大块片上工作区 SBUF 28 MiB UB 192 KB(+ L1 512 KB) 约 150×
矩阵累加器 PSUM 2 MiB L0C 128 KB 16×
矩阵输入 直接从 SBUF 读 必须先 staging 到 L0A/L0B,且要求分形格式 多两级缓冲与一级格式约束

Trainium2 讲 NeuronCore 的那两段话在昇腾同样成立,而且更严格。Trainium2 说片上存储是软件管理的,用它的原话说就是"从正确性角度,cache 并不存在"。昇腾上 UB / L1 / L0 之间没有任何自动搬运,每一步都必须写 DataCopy。唯一的例外是 L2 Cache:它有硬件管理策略,但只作用于 GM 访问路径,编程时不可见。

Cache 容量查询问题

环境上查容量有两个口径:

  • 运行时 GetCoreMemSize(CoreMemType::UB / L1 / L0_A / L0_B / L0_C),
  • 或读 $ASCEND_HOME_PATH/<arch>-linux/data/platform_config/<SoC>.ini。

但两个口径在 UB 与 L1 上并不一致(A2/910B 与 A3/910C 实测相同):

Buffer ini(配置值) GetCoreMemSize()(运行时) 差
L0_A / L0_B 65536 65536 0
L0_C 131072 131072 0
L1 524288 524032 −256
UB 196608 196352 −256
FB / BT 2048 / 1024 2048 / 1024 0

为什么 GetCoreMemSize() 比 ini 少 256 B?——因为它根本不是"物理容量"。 官方文档说得很清楚(ReserveLocalMemory):

该函数用于在 Unified Buffer 中预留指定大小的内存空间。调用该接口后,使用 GetCoreMemSize 可以获取实际可用的剩余 Unified Buffer 空间大小。

也就是说 GetCoreMemSize() 返回的是"当前还能用多少"(一个会随预留而变化的运行态数值),而不是芯片标称容量。而 kernel 里 TPipe::InitBuffer 的分配器按芯片总量分配 —— 这就是"两边数字对不上、但 196608 B 照样能分配满"的原因。

在910B4 卡上实测把 tile 大小不断翻倍(n = 2 × tile,每档都用 --verify 验正确性):

for T in 32736 32752 32768 32784 33024 33792 34816 40960; do echo "===== tile=$T  需要 $((3*T*2)) B"; python3 run_benchmark.py --kernel tiled -n $((T*2)) --tile $T --verify; done
tile 三个队列合计 结果
32736 196416 B ✅ 通过
32752 196512 B ✅ 通过
32768 196608 B ✅ 通过(结果正确)
32784 196704 B ❌ aclError:507035
33024 ~ 40960 198144 ~ 245760 B ❌ 全部失败

断点精确落在 196608 B 与 196704 B 之间,与 platform_config 的 ub_size=196608 完全吻合。

那 256 B 是什么?——已实测定位为"平台库的固定扣减",与预留无关。 探针里做了一次对照实验(预留 8 KiB 前后各取一次 UB):

[预留实验] 预留 8 KiB 前 UB = 196352,预留后 UB = 188160,差值 = 8192
           同一时刻 L1 = 524032(ReserveLocalMemory 只作用于 UB,看 L1 会不会跟着变)

差值恰好 8192(而不是 8192+256)→ 那 256 B 不是默认预留;L1 在预留后纹丝不动 → 也印证了该 API 只作用于 UB。

于是三者的账目关系就是:

数值(A2/910B4 实测) 说明
platform_config 的 ub_size 196608 芯片物理容量
GetCoreMemSize(UB) 196352 = 物理容量 − 256 B(平台库固定扣减)
再调用 ReserveLocalMemory(8K) 后 188160 再减 8192,与上面叠加
TPipe::InitBuffer 实际可分配 196608 分配器按物理总量算,不认那 256 B

L1 同样是"配置值 − 256"(524032 vs 524288),而 L0A/L0B/L0C/FB/BT 一分不差。至于那 256 B 到底代表什么(某个状态/元数据区?),公开资料里查不到——官方 ReservedSize 只有 8K/16K/32K 三档,没有它这一档。

那这份"预留"对 kernel 到底有没有约束力?——没有。 固定 tile = 32768(kernel 需要 196608 B),只改 Host 侧声明的预留:

python3 run_benchmark.py --kernel tiled -n 65536 --tile 32768 --reserve none --verify
python3 run_benchmark.py --kernel tiled -n 65536 --tile 32768 --reserve 8k   --verify
python3 run_benchmark.py --kernel tiled -n 65536 --tile 32768 --reserve 32k  --verify
--reserve GetCoreMemSize(UB) kernel 实际分配 结果
none(未声明预留) 196352 B 196608 B ✅ [ OK ]
8k 188160 B(−8192) 196608 B ✅ [ OK ]
32k 163584 B(−32768) 196608 B(比账本多 33024 B) ✅ [ OK ]

声明预留 32 KiB 之后,Tiling 侧的账本只剩 163584 B、比 kernel 实际要用的量少 33 KB,kernel 依旧分配成功、结果依旧正确;同时 L1 始终是 524032 B(不受影响,因为该 API 只作用于 UB)。这就是 --reserve 这个实验开关存在的意义。

于是整件事可以画成"两本互不约束的账":

┌─ Tiling(Host)侧 ────────────────────────────────┐
│  GetCoreMemSize(UB) = 物理容量 − 256B − 你的预留     │   ← 只是"报数"
│  196608      − 256  − {0 | 8192 | 32768}           │
│  = 196352 / 188160 / 163584                        │
└────────────────────────────────────────────────────┘
                ↕ 没有任何强制关系
┌─ kernel(Device)侧 ──────────────────────────────┐
│  TPipe::InitBuffer  →  只受物理容量 196608 B 约束   │   ← 真正的硬限制
└────────────────────────────────────────────────────┘

⚠️ 实践中必须注意:声明了预留 ≠ kernel 会让出空间。若你在 Tiling 里 ReserveLocalMemory()(或用了需要 workspace 的高阶 API,如 Matmul),就必须把这份扣减真正做进 tiling 计算(把算出来的更小 tile 传给 kernel)。否则 kernel 会照旧吃满 UB、与"留给别人的空间"重叠,而且不会报任何错,只表现为偶发数据错误。

实践结论:

  1. 容量预算用 ini / ub_size;把 GetCoreMemSize() 当作"扣除预留后的可用量"(偏保守);真卡边界时以 InitBuffer 的返回值 + 上机实测为准。
  2. 用了 ReserveLocalMemory() 就必须自己把扣减做进 tiling —— 它不会自动约束 kernel(见下)。

顺带一提,Ascend910B4.ini 里这几个键直接印证了上文的结构描述:ubbank_size=4096(= 每个 bank 4 KiB,与"192 KiB ÷ 48 bank"一致)、ubblock_size=32(= DataBlock 32 B)、vec_calc_size=128、cube_m/n/k_size=16。

各级存储的带宽测量

注意:有环境级查询接口 —— GetCoreMemBw()(文档):

void GetCoreMemBw(const CoreMemType &memType, uint64_t &bwSize) const;  // 单位 Byte/cycle
  • 官方只支持 L2 与 HBM 两个类型(L0_A/L0_B/L0_C/L1/UB 标注"预留参数,暂不支持");
atomgit@6ac4a72a17ae8e63171149c9:~$ bash ascend_probe_l0c.sh --cpp Ascend910B4
......
--- 运行结果 soc_version=Ascend910B4 ---
Buffer            Bytes            KiB
------------------------------------------
L0_A              65536           64.0
L0_B              65536           64.0
L0_C             131072          128.0
L1               524032          511.8
L2            100663296        98304.0
UB               196352          191.8
HBM         34359738368     33554432.0
FB                 2048            2.0
BT                 1024            1.0

[预留实验] 预留 8 KiB 前 UB = 196352,预留后 UB = 188160,差值 = 8192
           同一时刻 L1 = 524032(ReserveLocalMemory 只作用于 UB,看 L1 会不会跟着变)

GetCoreMemBw: L2 = 110 B/cycle, HBM = 32 B/cycle

AIC=20  AIV=40

对照上面的 §2:AIC 应等于 ai_core_cnt,AIV 应等于 vector_core_cnt。
若 AIC 只有 ai_core_cnt 的一半,说明一次只调度到一个 die。
  • 同一份数据在 platform_config/<SoC>.ini 的 [AICoreMemoryRates] 段(A2/910B4 实测):
atomgit@6ac4a72a17ae8e63171149c9:~$ grep -inE "bw|bandwidth|rate|pll" $ASCEND_HOME_PATH/aarch64-linux/data/platform_config/Ascend910B4.ini

[AICoreMemoryRates]
ddr_rate=32          ddr_read_rate=32    ddr_write_rate=32    // HBM(读/写同 32 B/cycle)
l2_rate=110          l2_read_rate=110    l2_write_rate=86     // L2
l1_to_l0_a_rate=512  l1_to_l0_b_rate=256 l1_to_ub_rate=128    // L1 → 各计算缓冲
l0_c_to_ub_rate=256  ub_to_l2_rate=64    ub_to_ddr_rate=64    // 回写路径
ub_to_l1_rate=128

换算成 GB/s:Byte/cycle × 实际时钟。可以在从profile的产物 OpBasicInfo.csv 的 Current Freq / Rated Freq 两列(dump/op_basic_info.txt 里同名同值)中获取。假设采样为 1650 MHz。以它计:HBM ≈ 32 × 1.65 ≈ 53 GB/s、L2 ≈ 110 × 1.65 ≈ 182 GB/s;用 1.8 GHz 计 L2 ≈ 198 GB/s。

边界:官方没写这些 rate 对应哪个时钟(核频/独立内存时钟)、是否乘引擎数,所以绝对 GB/s 只能给"×核频"的工程近似;相对用途(各通路带宽比例、瓶颈判断)则完全可信。GetCoreMemBw 只给 L2/HBM 两个数,L1/L0/UB 的 rate 只有 ini 里有。

L2 的"实测利用率"仍需 msprof op(usage_rate 列),L2 命中率在 L2Cache.csv。判断"带宽是否打满"用 usage_rate + 实测带宽就够,不必依赖未公开的峰值。SoC 传输带宽(Stars Soc Info)只出现在应用级导出(msprof --export)的 msprof_*.json timeline 里,算子级 CSV 没有。

结论(t512):UB 读写利用率 3.4% / 1.9%、L2 命中 1.7% —— 内存带宽远未打满,佐证瓶颈在指令/标量开销而非访存(详见报告 §1.3)。

避免 Unified Buffer 的 bank 冲突(A2 / A3)

UB = 192 KB 连续线性地址空间,按 32 B 切成 6144 个 DataBlock。软件端线性地址视角:

UB = 192 KB 连续线性地址空间,按 32 B 切成 6144 个 DataBlock

addr 0x00000 ─┬─ DataBlock 0    (32 B)
              ├─ DataBlock 1    (32 B)
              ├─ DataBlock 2    (32 B)
              │     ...
              ├─ DataBlock i    (32 B)   ← 任意 32B 对齐的地址 = 1 个 DataBlock
              │     ...
              └─ DataBlock 6143 (32 B)
                                  6144 × 32 B = 196608 B = 192 KB ✓

一次 repeat(向量指令)= 连续 8 个 DataBlock = 256 B
┌────┬────┬────┬────┬────┬────┬────┬────┐
│DB0 │DB1 │DB2 │DB3 │DB4 │DB5 │DB6 │DB7 │  = 256 B
└────┴────┴────┴────┴────┴────┴────┴────┘
  half:  8×16 = 128 个元素
  float: 8×8  = 64  个元素

UB 是 AIV 侧唯一的片上工作区,也是本作业 Part 1、Part 2 的主战场。UB 内部按 32 B 的粒度切分成了 bank,划分规则如下:

  • UB 总大小 192 KB,划分为 48 个 bank;
  • 每个 bank 由 128 行组成,每行 32 B(即每个 bank 4 KB);
  • 48 个 bank 进一步组织为 16 个 bank group,每组 3 个 bank,第 g 组由 bank g、bank g+16、bank g+32 组成(例如 bank15、bank31、bank47 同组);
  • Vector 单元每拍(一个指令周期)能从每个 bank group 读取或写入一行。

矢量计算通过Vector计算单元完成,矢量计算的源操作数和目的操作数均通过Unified Buffer(UB)来进行存储。Vector计算单元每个迭代会从UB中取出8个datablock(每个datablock数据块内部地址连续,长度32Byte),进行计算,并写入对应的8个datablock中。所以当前数据类型单次迭代时能处理的元素个数最大值为:256 / sizeof(数据类型)。

口径提示(实测发现,如实并列):platform_config/<SoC>.ini 里写的是 ubbank_num = 64、ubbank_size = 4096(64 × 4 KiB = 256 KiB,与同一份文件里 ub_size = 196608 的 192 KiB 对不上),而官方《避免 bank 冲突》文档是"192 KB 划分为 48 个 bank、16 个 bank group、第 g 组 = {g, g+16, g+32}"。

本节采用官方文档口径的 48,理由有两条:① 它是 bank 冲突优化的官方指导口径;② 它复现了官方那两个算例(blk_stride = 16 → 8 拍、blk_stride = 8 → 4 拍)。

值得说明的是,这两个算例无法区分实际的bank数量是 48 还是 64:决定冲突的是 块号 mod 16(bank group 序号),而 48 与 64 都是 16 的倍数,两种 bank 数下 8 个 DataBlock 的分组结果完全相同(blk_stride=16 → 全部落 group 0;blk_stride=8 → 两两落 group 0/8)。至于两者是否分别对应"可用 192 KiB"与"物理 256 KiB",公开资料没有说明,此处只并列实测值,不做推断。

图 4 UB 的 bank 结构,以及"8 个 DataBlock 落点"造成的拍数差异

当多条操作试图同时访问同一个 bank 或同一个 bank group 时,硬件只能让它们排队,于是本该一拍完成的操作被拆成多拍,这就是 bank 冲突。冲突分三类:

冲突类型 触发条件
读写冲突 同一条 Vector 指令的 src 与 dst 落在同一个 bank
写写冲突 dst 的多个 DataBlock 同时写同一个 bank group
读读冲突 多个 src(或同一 src 的多个 DataBlock)同时读同一个 bank group

官方文档给出的量化例子很直观:一个 Repeat 里 8 个 DataBlock(block0–block7),

  • blk_stride = 16 → 8 个 DataBlock 全部落在同一个 bank group → 全部冲突,8 拍才完成一个 Repeat;
  • blk_stride = 8 → 两两落在同一 bank group → 只有部分冲突,4 拍完成,快一倍。

两种修法:

① 优化计算逻辑:

把"跳读 + 连续写"改成"连续读 + 跳写"。

官方例子是对 (8, 16, 16) 做 (1, 0, 2) 的 transpose(输出 (16, 8, 16)):原实现中同一 Repeat 的 8 个输入 DataBlock 因 srcBlkStride = 16 全落在同一 bank group 而读读冲突;改成 srcBlkStride = 1、dstBlkStride = 8 后,同一 Repeat 内的 DataBlock 分散到不同 bank group,冲突消失。

一个地址是被两级译码的:

地址 → (bank 号, bank 内的行号 0..127, 行内偏移 0..31)

访存时候按 32 B 粒度交错(interleave):相邻的 32 B 落在不同 bank 上,48 个 bank 轮完一圈再从下一个 bank 的下一行继续。

输入是 (8,16,16) 的 half,所以:

量 值
一个元素 2 B
一个 DataBlock 32 B = 16 个元素
一整行(最内层 16 个元素) 16 × 2 B = 32 B = 正好 1 个 DataBlock = 一个 bank 的一行

于是这个张量一共 16 × 8 = 128 行,也就是需要128 个 DataBlock,行号(块号)= d0 × 16 + d1(d0 ∈ [0,8),d1 ∈ [0,16))。

向量Adds API 参考:高维切分API

// ① tensor 前 n 个数据计算(count 形式)
template <typename T, bool isSetMask = true>
__aicore__ inline void Adds(const LocalTensor<T>& dst, const LocalTensor<T>& src,
                            const T& scalarValue, const int32_t& count);

// ② tensor 高维切分计算 —— mask 数组模式
Adds(dst, src, scalarValue, uint64_t mask[], const uint8_t repeatTime, const UnaryRepeatParams& repeatParams);// dst = src + scalarValue    ← 一元:1 张量 + 1 标量

// ③ tensor 高维切分计算 —— mask 连续模式
Adds(dst, src, scalarValue, uint64_t mask,  const uint8_t repeatTime, const UnaryRepeatParams& repeatParams);

Add (dstLocal, src0Local, src1Local, ...)    // dst = src0 + src1          ← 二元:2 张量
  • mask 连续模式:表示前面连续的多少个元素参与计算。当操作数的数据类型占bit位16位时(如half/uint16_t),mask∈[1, 128];当操作数为32位时(如float/int32_t),mask∈[1, 64]。
  • mask 逐bit模式(数组模式):可以按位控制哪些元素参与计算,bit位的值为1表示参与计算,0表示不参与。当操作数为16位时,数组长度为2,mask[0]、$mask[1]∈[0, 2^{64}-1]$并且不同时为0;当操作数为32位时,数组长度为1,$mask[0]∈(0, 2^{64}-1]$;当操作数为64位时,数组长度为1,$mask[0]∈(0, 2^{32}-1]$。
数据类型 一次迭代最多元素 需要的 bit 数 mask 数组
half / int16 128 128 = 2×64 uint64_t mask[2]
float / int32 64 64 uint64_t mask[1]
int64 32 32(只用低 32 位) uint64_t mask[1](∈(0, 2³²−1])

跳读,连续写

举一个使用例子:

// int16_t数据类型单次迭代能处理的元素个数最大值为256/sizeof(int16_t) = 128,mask = 64,mask∈[1, 128],所以是合法输入
// repeatTime = 1, 共128个元素,单次迭代能处理128个元素,故repeatTime = 1
// dstBlkStride, src0BlkStride, src1BlkStride = 1, 单次迭代内连续读取和写入数据
// dstRepStride, src0RepStride, src1RepStride = 8, 迭代间的数据连续读取和写入
uint64_t mask = 64;
AscendC::Add(dstLocal, src0Local, src1Local, mask, 1, { 1, 1, 1, 8, 8, 8 });

默认参数,查看官方库grep -rn -A12 "struct UnaryRepeatParams" $ASCEND_HOME_PATH/aarch64-linux/asc/include/ 2>/dev/null | head -20:

// ascend 官方库自己的用法(matmul_client.h:1994-1995)
intriParams.dstBlkStride = (BLOCK_CUBE * sizeof(T) / DEFAULT_C0_SIZE);
intriParams.srcBlkStride = calcWidth * (BLOCK_CUBE * sizeof(T) / DEFAULT_C0_SIZE);
intriParams.dstRepStride = intriParams.dstBlkStride * DEFAULT_BLK_NUM;   // DEFAULT_BLK_NUM = 8(每迭代 8 块)
intriParams.srcRepStride = intriParams.srcBlkStride * DEFAULT_BLK_NUM;

原始实现:跳读,连续写。同一Repeat内输入的8个DataBlock都在同一个bank group而发生读读冲突。

跳读,连续写

示例代码:

uint64_t mask = 128;
UnaryRepeatParams params;
params.dstBlkStride  = 1;
params.srcBlkStride = 16;
for(uint32_t i=0; i<16; i++)   {
    AscendC::Adds(dstLocal[i * 128], srcLocal[i * 16], 0, mask, 1, params);
}

每次循环内部(同一 Repeat 内),读左图的一整列,写右图的一整行。i = 0(图上那个箭头标的就是它):

  • 读:in[d0 = 0..7][d1 = 0][:] → 块号 0, 16, 32, …, 112 → 左图第 0 列的那 8 格,竖着往下取,格间跳 16 块;
  • 写:out[d1 = 0][d0 = 0..7][:] → 块号 0, 8, 16, …, 56 → 右图第 0 行的那 8 格,横着连续写。

关键:LocalTensor[i] 的 i 是元素下标,不是块下标。所以块下标转为元素下标要 ×16。 代码里下标128、16、256,全都是「推进的块数 × 16」。

代码里的写法 换算成 DataBlock 为什么正好是这个数
原版 dst:dstLocal[i * 128] = 第 i × 8 个块的起点 一条指令连续写 8 块(repeatTime = 1 × 8 块),下一条指令要往后推 8 块;8 块 × 16 个元素/块 = 128
原版 src:srcLocal[i * 16] = 第 i 个块的起点 第 i 条指令从第 i 块开始(它读的是 d1 = i 的那些行),每条指令只推进 1 块;1 × 16 = 16
优化版 src:srcLocal[i * 256] = 第 i × 16 个块的起点 一条指令读 2 个迭代 × 8 个连续块 = 16 块,所以下一条要推进 16 块;16 × 16 = 256
优化版 dst:dstLocal[i * 16] = 第 i 个块的起点 目的侧改成跳写,第 i 条指令负责落点同余 i 的那些块,每条指令推进 1 块;1 × 16 = 16

优化实现: 连续读,跳写。同一个Repeat内输入的8个DataBlock不在同一个bank group内,避免了读读冲突。

连续读,跳写

示例代码:

uint64_t mask = 128;
UnaryRepeatParams params;
params.dstBlkStride  = 8;
params.srcBlkStride = 1;
for(uint32_t i=0; i<8; i++)   {
    AscendC::Adds(dstLocal[i * 16], srcLocal[i * 256], 0, mask, 2, params);
}

repeat 约束:矢量指令自身"一次调用最多能处理多少数据"的限制。

具体有三条约束:

  1. 每次迭代固定 8 个 DataBlock。 这不是你能调的参数,是硬件一次取数的宽度:blockNumber 在 UnaryRepeatParams 里是保留字段,默认就是 8。官方原话:"Vector 计算单元每个迭代会从 UB 中取出 8 个 datablock(每个 datablock 数据块内部地址连续,长度 32 Byte)"。

  2. mask 决定这次迭代里几个元素真的参与计算。 连续模式下取值范围是 256 / sizeof(T):half 最多 128、float 最多 64。这也是为什么官方样例里 half 用 mask = 128——正好把 256 B 铺满。

  3. repeatTime ≤ 255。 官方在高维切分那一节明确写了"由于硬件限制,repeatTime 不能超过 255"。

把 ①②③ 乘起来就是 单条指令的处理量上限 ≈ 63.75 KB。

优化版每个 chunk 要处理 16 个 datablock(= 2×8 块 = 256 个 half),一次迭代只能取 8 个 datablock(硬件规则:"每个迭代会从 UB 中取出 8 个 datablock…进行计算,并写入对应的 8 个 datablock"),所以这里 repeatTime = 2 。

另外还有一条容易忽略的:srcRepStride / dstRepStride 的类型是 uint8_t,所以"相邻迭代同序号块的跨距"最大 255 个 DataBlock。如果业务需要一次跳到 255 块(8160 B)以外,单条指令就表达不了,必须外面再套一层循环重新算基地址。

② 优化地址分配:

分配时适当扩大申请,把不同 Tensor 推到不同的 bank。官方例子是 z = x + y(各 4096 个 float):不对齐时 x、y 在同一 Repeat 内同读一个 bank group,且 x/y 与 z 读写同一个 bank;

// 原始:三个 Tensor 各 16 KB,落点相互撞车
pipe.InitBuffer(inQueueX, 1, 4096 * sizeof(float));
pipe.InitBuffer(inQueueY, 1, 4096 * sizeof(float));
pipe.InitBuffer(outQueueZ, 1, 4096 * sizeof(float));

// 优化:x 多申请 256 B,把 y 推开,避免一个 Repeat 内 x、y 同组;
//       y 多申请空间,保证 z 不与 x/y 落在同一个 bank
pipe.InitBuffer(inQueueX, 1, 4096 * sizeof(float) + 256);
pipe.InitBuffer(inQueueY, 1, 64 * 1024 - (4096 * sizeof(float) + 256));
pipe.InitBuffer(outQueueZ, 1, 4096 * sizeof(float));

原始实现,不做地址优化,直接使用InitBuffer分配内存,各个Tensor的地址分别为:

  • x:起始地址0x0,tensor长度为4096 * sizeof(float)字节

  • y:起始地址0x4000,tensor长度为4096 * sizeof(float)字节

  • z:起始地址0x8000,tensor长度为4096 * sizeof(float)字节

每个 tensor 是 4096 × 4 B = 16384 B = 512 个 DataBlock(32B)。这 512 个块不是堆在 4 个 bank 里,而是轮流落在 48 个 bank 上:

512 = 48 × 10 + 32
→ bank 0 ~ 31 各拿到 11 行,bank 32 ~ 47 各拿到 10 行

在一个Repeat内,x和y同时读同一个bank group,x/y和z同时读写同一个bank。

不做地址优化

优化实现,使用InitBuffer分配内存时适当扩大内存申请,各个Tensor的地址分别为:

  • x:起始地址0x0,tensor长度为(4096 * sizeof(float) + 256)字节

  • y:起始地址0x4100,tensor长度为(64 1024 - (4096 sizeof(float) + 256))字节

  • z:起始地址0x10000,tensor长度为4096 * sizeof(float) 字节

注:64 * 1024 = 65536 B = 64 KB,64 * 1024 是 16 个 bank 的空间。这里想让 z 落在 64 KB,x 占掉 16640 B,所以 y 只能是 65536 − (4096 × 4 + 256) B。

x多申请256字节,256 B = 8 块,刚好让 y 相对 x 的组号偏 8 → 跳出同组;避免一个Repeat内x和y同时读同一个bank group;

y多申请空间,确保z不会和x/y落入同一个bank

做地址优化

同组判据 同 bank 判据
两个块号相差是 16 的倍数 → 同一个 bank group 两个块号相差是 48 的倍数 → 同一个 bank

在本作业里意味着什么:

  • Part 1(vector add):顺序访问、blk_stride = 1 时一般不会撞冲突;但矩阵转置这类跳读/跳写的搬运与计算要按上面的第 ① 条设计访问模式,否则会比预期慢。
  • Part 2(fused conv + maxpool):UB 里通常同时存在输入 tile、权重 tile、im2col 结果和输出 tile。分配顺序与 padding 会直接决定拍数,建议从第一次实现起就按 256 B 粒度对齐分配,而不是等到调优阶段再补。
  • 观测手段:msprof 可采集"资源冲突占比"相关性能数据(对应 Trainium2 侧用 neuron-profile 看 timeline 的目的)。

Part 1: Learning the Ascend C Kernel Interface with Vector Add and Matrix Transpose (30 points)

在本节中,我们通过向量加的多个实现来介绍昇腾的编程模型,然后写一个简单的矩阵转置 kernel。

Trainium2 的 Part 1 代码在 /part1 目录下,是 Python 写的 NKI kernel + 一个 run_benchmark.py 命令行脚本。昇腾侧的对等实现已经写好,放在 part1/:kernel 是 kernels.cpp,Host 驱动是 main.cpp,命令行与 Trainium2 版一致:

 python3 run_benchmark.py --kernel {naive,tiled,stream,transpose} -n N [-m M] [--profile_name NAME]
  • 按 Step 1 / 2a / 2b 的要求改 kernels.cpp,并用 --tile(单次搬运元素数)、--block-dim(用几个 AI Core)扫参数——注意 matrix_transpose 目前是直通搬运骨架,Step 3 要你补上真正的转置;想看更正式的多文件算子工程结构,可以照官方样例(op_host/ + op_kernel/,或用 .asc 单文件异编译 + <<<>>> 直调);
  • 计时与剖析统一走 msprof(详见 Step 2b),或自建 host 侧计时循环(aclrtSynchronizeStream 前后取时间);
  • 练习入口可用 cann-learning-hub 与 cann/asc-devkit 的 examples/。

与 Trainium2 一一对应的三个知识点(分块并行、流式搬运、剖析驱动调优)在昇腾上全部成立,只是每一项的"旋钮"不同:Trainium2 拧 ROW_CHUNK / FREE_DIM,昇腾拧 tileLength(单次搬运字节数)与 BUFFER_NUM(双缓冲深度)。

注:关于kernel的调用方式,官方样例工程用目录名后缀区分开:*_frameworklaunch 是框架方式(op_host/ + op_kernel/),*_kernellaunch 是 kernel_name<<<numBlocks, l2ctrl, stream>>>(argument list);直调:

  1. numBlocks 是"逻辑核",不是"物理核" 每个 block 跑在某个物理核上、拿到逻辑 ID block_idx(GetBlockIdx()),一个物理核可以承载多个 block。官方建议"一般设置为物理核核数或其倍数",取值 [1, 65535]。对我们 A2/A3 单 die = 20 个 AI Core —— --block-dim 1 就是 numBlocks = 1。

  2. l2ctrl是 保留参数,官方原话:“暂时设置为固定值nullptr,开发者无需关注”。

  3. stream是执行流:维护异步执行顺序;核函数调用立即返回 host,真正执行排队在 stream 上;要拿结果/计时前必须 aclrtSynchronizeStream(stream)。

    实测时在 launch() 之后立刻 CHECK_ACL(aclrtSynchronizeStream(stream)) 再读时间——测到的就是"下发 + 执行 + 同步"全程。

用途 链接
官方样例总库(推荐入口,含 operator/ascendc/ 全部样例) gitee.com/ascend/samples
向量加 · Kernel 直调(单文件 add_custom.cpp + main.cpp 里 <<<>>> + CMakeLists + run.sh) 21_vectoradd_kernellaunch/VectorAddSingleCore
向量加 · 框架方式(op_host/ + op_kernel/ + framework/) 1_add_frameworklaunch/AddCustom
融合算子骨架(matmul+leakyrelu,对应 Part 2)直调 / 框架两版 13_matmulleakyrelu_kernellaunch · 12_matmulleakyrelu_frameworklaunch
矩阵乘双版本 / HelloWorld / 未对齐尾块样例 0_introduction
bank 冲突最佳实践(前面讨论的 x/y/z 地址分配、blkStride 转置) 4_best_practices/4_bank_conflict
CANN 新版样例仓库(asc-devkit,适配文档推荐) gitcode.com/cann/asc-devkit
文档:创建算子工程(msOpGen 生成 op_host/op_kernel) atlas_ascendc_10_0060
文档:简易自定义算子工程(上面的简化版) atlas_ascendc_10_0101
文档:基于样例工程完成 Kernel 直调 atlas_ascendc_10_0056

Ascend C Programming Model:

Ascend C 是华为面向算子开发场景的编程语言,原生支持 C/C++ 规范,编译产物通过 CANN 运行在昇腾 AI 处理器上。与 NKI 把 kernel 写成 Python 不同,Ascend C 的 kernel 是用 __global__ __aicore__ 标注的 C++ 函数,地位相当于 CUDA 的 __global__ 函数。

和 NKI 一样,Ascend C kernel 也只有三类操作:

  1. 搬运数据(CopyIn):从 GM 搬到片上 UB;
  2. 计算(Compute):在 Vector / Cube 计算单元上执行;
  3. 写回输出(CopyOut):从 UB 搬回 GM。

作为例子,下面用 Ascend C 实现同样的向量加。NKI的 @nki.jit 对应这里的 extern "C" __global__ __aicore__;NKI的 nl.ndarray(..., buffer=nl.hbm) 对应 GlobalTensor,nl.ndarray(..., buffer=nl.sbuf) 对应队列里的 LocalTensor。

#include "kernel_operator.h"
using namespace AscendC;

constexpr int32_t BUFFER_NUM = 2;   // 双缓冲深度

class KernelAdd {
public:
    __aicore__ inline KernelAdd() {}

    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z,
                                uint32_t totalLength, uint32_t tileLength) {
        this->tileLength = tileLength;
        this->tileNum = totalLength / tileLength;   // 本作业里 totalLength 保证整除
        // GM 上的输入输出(相当于 HBM 里的 NumPy 数组)
        xGm.SetGlobalBuffer((__gm__ half *)x, totalLength);
        yGm.SetGlobalBuffer((__gm__ half *)y, totalLength);
        zGm.SetGlobalBuffer((__gm__ half *)z, totalLength);
        // 片上 Buffer:BUFFER_NUM 份,这就是"双缓冲"
        pipe.InitBuffer(inQueueX, BUFFER_NUM, tileLength * sizeof(half));
        pipe.InitBuffer(inQueueY, BUFFER_NUM, tileLength * sizeof(half));
        pipe.InitBuffer(outQueueZ, BUFFER_NUM, tileLength * sizeof(half));
    }

    __aicore__ inline void Process() {
        // SPMD:每个 AI Core 只处理自己分到的那一段 tile(Trainium2 的 NKI 是单核,没有这一步)
        int32_t blockIdx = GetBlockIdx();
        int32_t blockNum = GetBlockNum();
        int32_t start = blockIdx * (tileNum / blockNum);
        int32_t end   = (blockIdx + 1) * (tileNum / blockNum);
        for (int32_t i = start; i < end; i++) {
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
    }

private:
    __aicore__ inline void CopyIn(int32_t progress) {
        LocalTensor<half> xLocal = inQueueX.AllocTensor<half>();
        LocalTensor<half> yLocal = inQueueY.AllocTensor<half>();
        DataCopy(xLocal, xGm[progress * tileLength], tileLength);   // GM -> UB
        DataCopy(yLocal, yGm[progress * tileLength], tileLength);
        inQueueX.EnQue(xLocal);                                     // 入队即完成同步
        inQueueY.EnQue(yLocal);
    }

    __aicore__ inline void Compute(int32_t progress) {
        LocalTensor<half> xLocal = inQueueX.DeQue<half>();
        LocalTensor<half> yLocal = inQueueY.DeQue<half>();
        LocalTensor<half> zLocal = outQueueZ.AllocTensor<half>();
        Add(zLocal, xLocal, yLocal, tileLength);                    // 逐元素向量加
        inQueueX.FreeTensor(xLocal);
        inQueueY.FreeTensor(yLocal);
        outQueueZ.EnQue<half>(zLocal);
    }

    __aicore__ inline void CopyOut(int32_t progress) {
        LocalTensor<half> zLocal = outQueueZ.DeQue<half>();
        DataCopy(zGm[progress * tileLength], zLocal, tileLength);   // UB -> GM
        outQueueZ.FreeTensor(zLocal);
    }

    TPipe pipe;  // 核内内存 / 队列管理器(管 UB 的分配)
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX, inQueueY; // 搬入队列,数据落在 UB;第二个模板参数 = 缓冲份数(2 = 双缓冲)
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueZ; // 搬出队列(也在 UB)
    GlobalTensor<half> xGm, yGm, zGm; // GM 上的张量,即设备侧地址的 C++ 包装
    uint32_t tileLength, tileNum; // 存下分块大小与块数
};

// 设备侧入口:一段 kernel 会被很多个 AI Core 同时执行(SPMD)
extern "C" __global__ __aicore__ void add_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z,
                                                 uint32_t totalLength, uint32_t tileLength) {
    KernelAdd op;
    op.Init(x, y, z, totalLength, tileLength);
    op.Process();
}

对照上面的代码,注意这几点:

  • x、y、z 是 GM 上的地址,用 GlobalTensor 包一层即可访问,相当于 HBM 里由 NumPy 创建的数组。
  • xLocal、yLocal 是 DataCopy 搬到 UB 之后的片上张量,由 TQue 分配,相当于 SBUF 里的 a_sbuf / b_sbuf。
  • Add(zLocal, xLocal, yLocal, tileLength) 用 Vector 单元做逐元素加法;如果要加一个常量标量,用 Adds。
  • DataCopy 在 GM 与片上 Buffer 之间搬数据,概念上等同 nisa.dma_copy(也等同 CUDA 的 cudaMemcpyAsync)。
  • EnQue / DeQue 不只是"放队列":它们由框架负责插入依赖同步,并配合 BUFFER_NUM = 2 实现双缓冲。NKI的 nisa.dma_copy 是显式、无队列的,这是两个编程模型最核心的体感差异。

关于"数据必须是二维 tile"这件事:注意 NKI 的片上张量是二维数组(partition dimension 最大 128,free dimension 最大 64K),并因此限制了 naive 版本只能处理不超过 128 个元素。昇腾没有这套二维语义:Vector 侧的 LocalTensor 就是一段一维片上内存,一次能处理多少元素由 UB 容量和指令的长度/repeat 约束决定。但昇腾有另一套同样硬的约束,写在 API 的约束描述里:

  • Vector 指令的源与目的都必须在 UB 中,首地址与操作长度通常要求 32 B 对齐;
  • 进入 Cube 的数据(Part 2 才用到)必须满足 16×16 分形与 L0A/L0B 的 512 B 对齐。

也就是说,"为什么 naive 版本做不大"这个问题在两个平台上都成立,只是物理原因不同。

Step 1: Chunking Vectors to Parallelize Across 128 Compute Lanes (6 points)

要让向量加能处理大于一次搬运量的数据,就要分块(把原张量切成子集,逐块处理)。

对比NeuronCores,昇腾没有"128 条 lane"这条固定说法。官方给的是按字节的宽度——一条矢量指令每次迭代取连续的 8 个 DataBlock = 256 B:16 位类型(half)刚好是 128 个元素(mask 上限 = 256 / sizeof(T)),32 位类型(float)则是 64 个。

注:platform_config 里有 vec_calc_size = 128 这个键。它不在 CANN 公开手册里(全量文档索引搜不到),但它的身份可以从 CANN 源码查到:

  • 谁读它:cann/runtime 的 src/platform/platform_info.cpp(8.5.0,L338)把 INI 键解析进 platform_info_temp.ai_core_spec.vec_calc_size;Vector Core 段在 L600 另有一份 vector_core_spec.vec_calc_size;
  • 它住哪:pkg_inc/platform/platform_info_def.h 的 tag_ai_core_spec,紧挨着 cube_m_size / cube_n_size / cube_k_size;
  • 在 INI 里的位置:[AICoreSpec] 段中 cube_m/n/k_size=16 之后、各级存储容量(l0_* / l1_size)之前。

于是它的含义很清楚:与"Cube 一次算 16×16×16"对称,vec_calc_size 是"Vector 单元一次计算能处理的元素个数"(以 16 位元素计)。三个数字互相吻合:ubblock_size = 32 B × 每迭代 8 个 DataBlock = 256 B,256 ÷ 2 B = 128 ✓;换成 32 位类型就是 64 —— 正是 mask ∈ [1, 64] 的由来。

所以它是"硬件宽度"(一拍能算多少元素),不是 Trainium 那种"数据布局约束"——下面这张表就是说清这个区别:

Trainium2(NKI) 昇腾(Ascend C)
那个 128 是什么 数据布局约束:SBUF / PSUM 里的张量必须是二维 (P, F),且 P ≤ 128 硬件宽度:一条指令每次迭代取 256 B = 128 个 half,与数据怎么摆放无关
还有没有"圈内的维度" 有 —— free 维 F,单条指令内每个 partition 最多 64K 个元素 没有。一次迭代的 128 条 lane 各处理 1 个 half 就满了
想处理更多元素 把 F 调大(同一条指令内并行) 只能 repeatTime(≤ 255)多次迭代(多个周期,靠流水重叠)

所以本步骤在昇腾上要关心的不是"哪一维必须 ≤ 128",而是单次搬运多少元素(tileLength)以及怎么让它与计算重叠;核内的并行度是固定的(每拍 128 个 half → mask → repeatTime),不是本步骤的旋钮。

// 在 Init 里把 tileLength 当作旋钮;Process 里按它循环
// TODO: 你需要为本步骤修改这个变量
uint32_t tileLength = 1;      // 单次搬运并计算的元素个数

Process() 里的循环相当于 Trainium2 的 nl.affine_range:块与块之间没有循环携带依赖(loop-carried dependency),所以可以放心地让框架做流水优化。若真有依赖,需要自己保证顺序(或用 SetFlag/WaitFlag)。

你需要做的事:

  1. 用 tileLength = 1 跑向量加,向量长度 25600。执行时间是多少微秒(μs)?

  2. 昇腾上没有"partition dim ≤ 128"这条限制,昇腾侧的限制来自 UB 容量。把 tileLength 逐步调大,直到接近 UB 能装下的上限(注意 BUFFER_NUM = 2 时每个张量要占 两份空间:192 KB 在三个张量之间均分,每个张量 64 KB,也就是每份缓冲 32 KB,fp16 下 tileLength 最多 16384)。记录 tileLength 取较大值时的执行时间,并回答:相比 tileLength = 1 快了多少?为什么更快?(提示: 一次搬运摊薄了 MTE2/MTE3 的 setup 开销,一条 SIMD 指令覆盖更多元素则摊薄了指令发射开销。)

  3. 继续把 tileLength 调大到超出 UB 容量,你应该会看到 InitBuffer 申请失败(或在运行/编译期触发对齐与长度断言)。用一句话解释为什么会报错。

Step 2a: Improved Data Streaming (4 points)

目前为止,我们用的分块方式只保证"能跑",还没有针对搬运效率做设计。

NKI 的做法是把一维向量 reshape 成 (128, FREE_DIM) 的二维 tile,让一次 dma_copy 搬更多数据,从而摊薄 DMA 传输的建立开销(每个 NeuronCore 有 16 个 DMA 引擎,但每次传输都要有 setup 成本)。

昇腾的 DataCopy 本来就是"一次搬 tileLength 个元素",没有分区维/自由维的二维语义。因此这一步在昇腾上被拆成两个独立的旋钮,需要分别测量:

// 旋钮 1:单次搬运的粒度
uint32_t tileLength = ...;      // 对应 Trainium2 的 FREE_DIM(越大,搬运次数越少)
// 旋钮 2:流水深度
constexpr int32_t BUFFER_NUM = 2;   // 1 = 无流水(搬完再算),2 = 双缓冲

你需要做的事:

  1. 固定总数据量,比较 BUFFER_NUM = 1 与 BUFFER_NUM = 2 的执行时间。快了多少?为什么?

  2. 固定 BUFFER_NUM = 2,把 tileLength 从很小调到接近 UB 上限。搬运次数(可由 msprof 看到)随 tileLength 增大而单调下降,但耗时不一定是单调下降的。找到你测得的最优点,并解释:为什么搬运次数更少反而可能更慢?(提示:往下一节和 Part 2 的"流水线重叠"方向想。)

Step 2b: Learning to Use msProf (5 points)

选择 tile 大小是一个权衡:

  1. tile 太小,指令开销显著,引擎执行效率低;
  2. tile 太大,引擎之间的流水重叠变差,且 UB 压力上升("UB 压力"指 UB 可能被占满,双缓冲退化为单缓冲)。

Trainium2 用 neuron-profile 的 GUI 观察 timeline。昇腾上的对应工具是 msprof(算子级 msprof op / 仿真级 msprof op simulator)与 MindStudio Insight(timeline 可视化),不需要 InfluxDB,也不需要 SSH 端口转发。

msprof 是主命令,后面跟不同的"模式词"决定采集类型。不带模式词 = 应用/模型级;带 op = 算子上板;带 op simulator = 算子仿真。

msprof [通用参数] <你的程序>              → 应用/模型级(PROF_XXX)
msprof op [算子参数] <你的程序>            → 算子上板(OPPROF_XXX)
msprof op simulator [仿真参数] <你的程序>  → 算子仿真(OPPROF_XXX/simulator/)

常用命令:

# 单算子,最简
msprof op ./add_custom_npu

# 指定输出 + 指标
msprof op --aic-metrics=Default --output=./output_data ./add_custom_npu

# 多算子:只采 Add 和 Sub,各采 10 次
msprof op --launch-count=10 --kernel-name="Add|Sub" --output=./output_data ./test

# 按 kernel 名精确采(Triton/PyTorch 场景常用)
msprof op --kernel-name=_layer_norm_fwd_fused --output=$HOME/projects/output \
          python3 03-layer-norm.py

# 指定设备
export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3
msprof op ./00_basic_matmul 256 512 1024 0

# 基于算子二进制 .o 的配置文件
msprof op --config=./add_test.json --aic-metrics=Default --output=./output_data

指标口径对照:

trainium 观察项(neuron-profile) 昇腾对应观察项(msprof / MindStudio Insight)
dma_transfer_count MTE2 / MTE3 的搬运次数与搬运量
total_time 算子总耗时(msprof op 输出)
DMA-E79 timeline MTE2 / MTE3 的 pipeline 视图
VectorE timeline Vector 单元的 pipeline 视图
Pending DMA Count 搬运队列堆积情况
DMA Throughput MTE 带宽利用率
(无直接对应) MTE2 / Vector / MTE3 的流水重叠度,昇腾上这是显式的,看得很清楚

两种 msprof 输出,两种导入法:

应用/模型级 msprof 算子级 msprof op
产出目录 PROF_XXX/ OPPROF_XXX/
关键文件 mindstudio_profiler_output/ 下的 msprof_*.json、op_summary_*.csv;或 msprof_*.db visualize_data.bin
Insight 导入什么 整个 PROF_XXX 目录(或它的上一层) visualize_data.bin 单个文件
展示页签 时间线 / 内存 / 算子 / 概览 / 通信 详情(Details) / 源码(Source) / 缓存(Cache) / 时间线(Timeline)
是否需要先 export 需要 msprof --export=on 解析 采集完直接可用

算子级 msprof op 场景

算子级MindStudio可视化操作(假设msprof op输出的文件夹为s2b_stream_t512):

  1. Insight 点左上角 导入数据 → 选择性能数据路径
  2. 直接选到 .../s2b_stream_t512/visualize_data.bin 这个文件
  3. 导入后去「详情(Details)」页签看,别盯着 Timeline

上板(msprof op 非 simulator)模式采集的 visualize_data.bin,Insight 会呈现这几张图:

  • 计算内存热力图:算子基础信息、计算工作负载分析、内存负载分析
  • Roofline 瓶颈分析图:判断是计算瓶颈还是内存瓶颈
  • Cache 热力图:L2 Cache 访问命中情况
  • 通算流水图、算子代码热点图

指令流水图(Timeline)只在这两种情况下有:

  • msopprof simulator 仿真采集 → 生成 trace.json(导入它才有指令流水 Timeline)
  • Atlas 350 加速卡 通过 msopprof 采集指令流水图

910B / A2 上跑 msprof op 上板采集,结果落在 Details,不在 Timeline。

指令级流水Timeline仿真采集:

msprof op simulator --soc-version=Ascend910B1 --output=./op_prof_sim ./your_op_bin

产出 trace.json + visualize_data.bin,前者导 Insight 看 Timeline 指令流水,后者看源码热点(Source 页签)。

算子级 msprof op产出的 csv文件:

CSV 含义
OpBasicInfo.csv 算子基础信息(耗时、核数等)
ArithmeticUtilization.csv Cube / Vector 指令 cycle 占比
PipeUtilization.csv 计算单元与搬运单元耗时占比(MTE2/VECTOR/MTE3 等)
Memory.csv UB、L1 与主存读写带宽速率
MemoryL0.csv L0A/L0B/L0C 读写带宽速率
MemoryUB.csv Vector / Scalar 到 UB 的读写带宽速率
L2Cache.csv L2 Cache 命中率
ResourceConflictRatio.csv bank group / bank conflict / 资源冲突占比

本部分主要查看最后三个:PipeUtilization.csv 看流水是否重叠(MTE2/V/MTE3 占比是否均衡),ResourceConflictRatio.csv 看 bank 冲突率,MemoryUB.csv 看 UB 带宽是否打满。

应用级 msprof 场景

要跑模型级 profiling,完整链路是:

# 1. 采集
msprof --application="python3 train.py" --output=./prof_data
# 2. 解析(必须!不 export 就没有 mindstudio_profiler_output)
msprof --export=on --output=./prof_data
# 3. Insight 导入 ./prof_data/PROF_XXX 这一层

框架数据的目录名有讲究:PyTorch 数据要是以 ascend_pt 结尾的目录,MindSpore 要 ascend_ms 结尾,且这两者需要放在"父目录/xxx_ascend_pt"这样的固定层级下,多套一层就会被误判成集群数据。集群场景若 repeat>1,要按时间戳拆成多份分目录导入。

文件大小建议:bin ≤ 500MB、json ≤ 1GB、csv ≤ 500MB,超了会明显变慢或失败。

应用/模型级 msprof --export 输出文件:mindstudio_profiler_output/msprof_*.json + op_summary_*.csv → Insight 的 Timeline / Operator / Summary 页签,回答"整网哪个算子慢、搬运 vs 计算、模型级 Roofline"

实测带宽与利用率

容量之外,性能分析还要看各级存储的实测带宽与利用率。msprof op 每个产出目录里都有专门的 CSV(下表实测值取自本作业 s2b_stream_t512/,n=25600、单核、纯 Vector 路径):

存储通路 文件 · 列 含义 实测(t512)
UB 读写(Vector 引擎视角) MemoryUB.csv · aiv_ub_read_bw_vector / aiv_ub_write_bw_vector Vector 读写 UB 的实测带宽 6.59 / 3.29 GB/s
GM ↔ UB(搬运引擎视角) Memory.csv · aiv_gm_to_ub_bw / aiv_ub_to_gm_bw MTE2 搬入 / MTE3 搬出的有效带宽 6.59 / 3.29 GB/s
主存(HBM) Memory.csv · aiv_main_mem_read_bw / aiv_main_mem_write_bw 对 Global Memory 的读写带宽 6.78 / 3.31 GB/s
L0A / L0B / L0C MemoryL0.csv · aic_l0x_*_bw Cube 侧 L0 读写带宽;纯 Vector 算子是 NA(走 Cube 路径才统计) NA
L1 Memory.csv · aic_l1_read_bw / aic_l1_write_bw Cube 侧 L1 带宽;纯 Vector 算子为 NA NA
L2 L2Cache.csv 只有命中率,没有带宽列 命中率 1.73%

搬运指令数也在 Memory.csv 里(aiv_mte2_instructions / aiv_mte3_instructions),本作业 --tile 512 为 101 / 51,正好是"50 轮 × 2 次搬入 +1"和"50 轮搬出 +1"。

利用率与"参考上限":Memory.csv 的 GM_to_UB_bw_usage_rate(%) / UB_to_GM_bw_usage_rate(%) 是 msprof 按某个参考带宽算的利用率。用"实测带宽 ÷ 利用率"可反推参考值:

GM→UB:6.59 GB/s ÷ 3.38% ≈ 195 GB/s
UB→GM:3.29 GB/s ÷ 1.87% ≈ 176 GB/s

你需要做的事:

  1. 用 msprof 采集你把 tileLength 调成"很大"与"较小"两组的指标。各自的耗时与搬运次数是多少?

  2. 从你测到的数据里总结耗时与搬运次数的关系,并解释它。(注意:在本环境里,大 tile 的搬运次数少、通常更快;但把 tile 拉到极小(如 512)时,搬运次数暴涨、耗时也暴涨——这正好与Trainium2 版的 FREE_DIM = 1000-vs-2000(搬运次数更多反而更快)相反。两边共用同一个 Hint:pipelining;差别在于对比落在了流水线的不同区间。)

    回答时请用你自己的实测数据:慢的到底发生在哪?可看 PipeUtilization.csv 的 scalar(标量) / Vector / MTE2 / MTE3 占比、Memory.csv 的搬运指令数与有效带宽、ResourceConflictRatio.csv 的等待占比;再做一组同 tile 的 BUFFER_NUM = 2 vs 1 对照,解释"搬运次数完全相同、双缓冲更快"。可附 MindStudio Insight 截图(或自绘引擎占用图),并配文字说明它如何验证你的结论。

  3. 另外看一眼 UB 的资源冲突占比(bank 冲突),确认你的 tileLength 与 LocalTensor 分配方式没有引入额外冲突(背景见 Part 0 的"避免 Unified Buffer 的 bank 冲突")。

Step 3: Matrix Transpose (15 points = 10 coding + 5 writeup (+ 1 EC))

Matrix Operations on an AI Core

开始之前,先看昇腾是怎么做矩阵运算的。前面说过,AI Core 里有多种计算单元,各自针对不同类型的算术操作做了优化。

NeuronCore 这里是 128×128 的脉动阵列(Tensor Engine),输入来自 SBUF、输出写进 PSUM。昇腾的对应物是 Cube 单元:一条指令完成两个 16×16 fp16 矩阵乘(分形 16×16×16),左矩阵在 L0A、右矩阵在 L0B、结果写进 L0C。三者都是片上存储,其中 L0C 是累加器,支持"读-加-写",因此适合分块累加的大矩阵乘:这一点与NeuronCore 讲 PSUM 的逻辑完全一致,只是容量从 2 MiB 变成了 128 KB(见 Part 0 的对照表)。

Writing the kernel

在进入 Part 2 真正的矩阵乘之前,先写一个"小 kernel":转置一个矩阵。你的 kernel 接收形状为 (M, N) 的二维张量,返回形状为 (N, M) 的张量。

extern "C" __global__ __aicore__ void matrix_transpose(GM_ADDR a, GM_ADDR out, ...) {
    // TODO: 你的实现。
    // 建议:分块 128×128(或按 UB 容量决定的块大小)搬运到 UB,
    //       调用转置接口后搬回 GM。
}

[!IMPORTANT]

这一步与 Trainium2 的差异最大,把两边的规定逐字对照着看。

Trainium2 版写的是:

To actually perform the transpose, you must call nisa.nc_transpose, which is a built-in instruction that uses the Tensor Engine to transpose tiles of size up to 128x128, storing the result in PSUM. You are not allowed to use other compute instructions, including nisa.dma_tranpose or nl.transpose. (Memory instructions, including nisa.dma_copy and nl.ndarray, are of course allowed.)

昇腾侧对应的规定是:

要对数据块做转置,你必须调用 AscendC::Transpose。它在 UB 内完成一个 16×16 数据块的转置,源与目的都是 LocalTensor(TPosition ∈ VECIN / VECCALC / VECOUT,起始地址 32 B 对齐),支持的数据类型为 half / uint16_t / int16_t。

当一次要转的块大于 16×16 时,用 AscendC::TransDataTo5HD(官方原话:"完成转置功能时,相比于 Transpose 接口,Transpose 仅支持 16×16 大小的矩阵")。

不允许把转置"外包"出 kernel——例如在 Host 侧用 torch.transpose(x, 0, 1).contiguous() 或 AclNN 的 Transpose 算子先转好再喂进来,也不允许用"先改成 NZ 格式、指望下游替我做转置"的方式绕过。搬运类指令(DataCopy、DataCopyPad)当然随便用。

两者逐项对照:

Trainium2 nisa.nc_transpose 昇腾 AscendC::Transpose / TransDataTo5HD
语义层级 Tensor Engine 的计算指令 搬运 / 格式转换类 API,不是 Cube 的计算指令
结果落点 PSUM(2 MiB 累加器) UB(src 与 dst 都是片上 LocalTensor)
单次粒度 单指令 ≤ 128×128 普通 Transpose 只支持 16×16;更大的块用 TransDataTo5HD
额外能力 — 增强版 Transpose 通过 TransposeParamsExt 还支持 [N,C,H,W] ↔ [N,H,W,C] 互转
对齐要求 PSUM 相关约束 起始地址 32 B 对齐(UB 通用约束)
"不许用什么" 禁 dma_transpose / nl.transpose 没有对应禁令(转置本就是搬运语义);等价规定是上面那条"不许外包给框架算子"

一个会改变实现结构的后果:Trainium2 一条 nc_transpose 能吞下 128×128,而你这里一次 Transpose 只有 16×16。所以同一个转置任务,在昇腾上要自己组织循环(对 UB 里的 tile 按 16×16 小块逐个转,或用 TransDataTo5HD 一次做更大的块),这也意味着分块策略对性能的影响比在 Trainium 上大得多。

进阶提示,留给 Part 2 用:如果转置只是为了喂给下一个矩阵乘,可以在搬入 L0A/L0B 时用 LoadDataWithTranspose 顺带转置,甚至直接用 Matmul 的转置参数把这次转置整体消掉——这对应 NKI 里的 lhsT 约定。但 Part 1 这一步请按上面的规定,老老实实用 Transpose 做出来。

你需要做的事:

  1. 完成 kernel,在 1024×1024 的矩阵上测试,记录执行时间(μs)。

  2. 不使用剖析工具,先判断你的 kernel 是 memory-bound 还是 compute-bound,并说明理由。然后用与 Step 2b 相同的方式做剖析验证(可以附截图,但请给出文字说明它是如何验证你的结论的)。

  3. (额外加分,1 分) 优化你的实现以最小化延迟。Trainium2 的达标线是 4096×4096 上 <700 μs —— 昇腾上没有现成阈值,你需要自建基线:

    • 基线 1:torch.transpose(x, 0, 1).contiguous() 在 NPU 上(即 CANN 自带实现)的耗时;
    • 基线 2:你自己的朴素分块实现。

      达标条件建议定为:打败基线 1,或达到其 1.2 倍以内。值得注意的是,昇腾上最有"平台味道"的优化方向是消除物理转置:如果转置只是后续算子的输入,直接在下一跳使用 Matmul 的 transposeA / transposeB 开关或换一种内存布局,比真的搬一遍快得多。

Part 2: Implementing a Fused Convolution - Max Pool Layer (70 points)

既然你已经学会在 AI Core 上高效搬数据,接下来就该真正写一个昇腾 kernel 了。本部分要求实现一个同时完成 convolution 与 max pooling 的 kernel —— 这两个操作是现代卷积神经网络(CNN)的基本组件。关键要求同样是融合:中间结果不能落到片外 HBM。

An Ascend C Matrix Multiplication Kernel

Trainium2 版在这里用一个分块的 NKI 矩阵乘做示范。昇腾上的对应物是 Cube 单元 + Matmul 高阶 API。

先看硬件约束的对应关系。 NKI 版给出的是 Tensor Engine 的三个 tile 上限(A ≤ 128×128、B ≤ 128×512、C 在 PSUM ≤ 128×512)——那是"脉动阵列的入口宽度"。昇腾的 Cube 没有这种形状约束,取而代之的是三块片上 Buffer 的容量上限:

Trainium2(Tensor Engine) 昇腾 A2 / A3(Cube 单元)
左矩阵 A ≤ (128, 128) L0A = 64 KB → fp16 下 baseM × baseK ≤ 32768 元素
右矩阵 B ≤ (128, 512) L0B = 64 KB → fp16 下 baseK × baseN ≤ 32768 元素
结果 C 在 PSUM,≤ (128, 512) L0C = 128 KB → fp16 下 baseM × baseN ≤ 65536 元素,fp32 下减半
计算粒度 128×128 脉动阵列 16×16×16 分形(cube_m/n/k_size = 16)

注意这两组约束的性质不同:Trainium2 是"形状上限",昇腾是"容量上限"。容量上限意味着 tile 的乘积受限、但单个维度可以更灵活——baseM 可以取 128 也可以取更小,只要 baseM × baseN 不超 L0C 容量即可。

所有容量的实测值与推导见 Part 0。这些值是架构级常数(910B 与 910C 完全相同),换机型不必重算 tile;但核数会变,所以要运行时取。

分块矩阵乘的骨架。 与 NKI 版一样,任意尺寸的矩阵乘在昇腾上也必须拆成"固定尺寸 tile 上的若干次矩阵乘"——因为 L0A / L0B / L0C 都装不下整个矩阵。昇腾提供两个层次:

  • Matmul 高阶 API(推荐作为起点):内部自动完成 GM → L1 → L0A/L0B 的 staging、L0C 上的累加、以及多核切分;你只需在 host 侧填好 TCubeTiling 结构体;
  • Mmad 基础 API:自己管理 L0A/L0B 的搬运与 L0C 的累加(差异见分离模式),只在需要极致控制时使用。

高阶 API 的分块骨架大致如下(完整可运行样例见 matmul_leakyrelu):

// 分块矩阵乘 + 逐块处理 C 的骨架(对应 NKI 版的 nki_matmul_tiled_)
template <typename aType, typename bType, typename cType, typename biasType>
__aicore__ inline void MatmulTiled(GM_ADDR a, GM_ADDR b, GM_ADDR bias, GM_ADDR c,
                                  GM_ADDR workspace, TCubeTiling &tiling)
{
    TPipe pipe;
    TQue<TPosition::VECOUT, 1> outQueue;
    pipe.InitBuffer(outQueue, 1, tiling.baseM * tiling.baseN * sizeof(cType));

    matmul::Matmul<matmul::MatmulType<TPosition::GM,    CubeFormat::ND, aType>,
                   matmul::MatmulType<TPosition::GM,    CubeFormat::ND, bType>,
                   matmul::MatmulType<TPosition::VECIN, CubeFormat::ND, cType>,   // C 直接落在 UB
                   matmul::MatmulType<TPosition::GM,    CubeFormat::ND, biasType>> mm;
    REGIST_MATMUL_OBJ(&pipe, GetSysWorkSpacePtr(), mm, &tiling);

    mm.SetTensorA(gm_a);
    mm.SetTensorB(gm_b);
    mm.SetBias(gm_bias);
    while (mm.template Iterate<true>()) {                      // 每次算出一块 baseM × baseN 的 C
        auto cLocal = outQueue.AllocTensor<cType>();
        mm.template GetTensorC<true>(cLocal, false, false);    // Cube → UB,交接由框架完成
        outQueue.EnQue(cLocal);
        /* TODO: 在这里处理这一块 C(加 bias、maxpool……) */
        outQueue.DeQue<cType>();
        outQueue.FreeTensor(cLocal);
    }
    mm.End();
}

三个关键对应关系:

  • Iterate() + GetTensorC() ↔ nisa.nc_matmul + nl.copy。

    • NKI 版里 res_psum += nisa.nc_matmul(...) 手动累加进 PSUM,循环结束后再 nl.copy(res_psum, ...) 搬回 SBUF;
    • 昇腾上 Iterate() 一次算出一块 baseM × baseN 的 C,GetTensorC(ubLocalTensor) 直接把它取到 UB(逻辑位置 CO2 → VECIN)。L0C 的累加由框架做,你不需要写 res +=。
  • TCubeTiling ↔ NKI 的 TILE_M / TILE_K / TILE_N。 NKI 里 tile 尺寸是 kernel 内的常量(128 / 128 / 512);昇腾的 tile 尺寸由 host 侧的 Tiling 函数算好、通过 TCubeTiling 传给 kernel。这是"Host/Kernel 分离"与"单文件 kernel"在工程结构上的差异,也是昇腾算子要写 host + kernel 两份代码的原因。

  • GetSysWorkSpacePtr() / GetLibApiWorkSpaceSize()。 Matmul 需要一块 workspace,其大小由库决定、与你算的矩阵规模无关。

把这条算法映射到昇腾,有六个点必须重新决策:

# 决策点 Trainium2 的做法 昇腾上你必须做的选择
1 用几个核 单 NeuronCore 本环境一个 die 给 20 个 AIC / 40 个 AIV(A3 的 24 核变体为 24/48)。blockDim 在 Host 侧运行时用 GetCoreNumAic() 取,不要硬编码
2 输出 tile 取多大 PSUM 装得下 128×512 L0C 只有 128 KB:fp32 下 baseM × baseN ≤ 32768 元素(baseM = 128 时 baseN ≤ 256),fp16 下 ≤ 65536。L0B 另有 baseK × baseN ≤ 32768(fp16)的约束,要与 L0C 联立求解
3 K 循环放在哪 K = C_in,分块累加进 PSUM 同样累加进 L0C,但每个 K 块多一级 MTE1: L1 → L0A/L0B 的搬运
4 maxpool 在哪做 nisa.tensor_reduce(nl.max, ...) on SBUF 在 UB 上做:ReduceMax / WholeReduceMax,或三次逐元素 Max
5 AIC 与 AIV 怎么配合 不适用(三种引擎同核) Cube 在 AIC、Vector 在 AIV,是两个核;本环境实测 1 个 AIC 配 2 个 AIV。bias 与 maxpool 属于 Vector 侧,跨核交接见 Fusing Convolution and Max Pool
6 权重能不能常驻片上 README 保证"权重总能完整放进 SBUF(28 MiB)" ⚠️ 这条保证在昇腾上不成立:fp16 下单个 (C_out, C_in) 切片 256×256×2 B = 128 KB,9 个共 1.125 MiB,而 L1 只有 512 KB。单切片放得下、全量放不下——必须按 (i,j) 从 GM 重新加载权重,或调整循环顺序让同一 (i,j) 的权重在整块 N 上复用

Convolution Layer Overview

数学定义与 Trainium2 完全一致,逐字成立,本节只重述要点以便自查。

卷积在输入 feature map 上滑动一个 filter:每个位置上,filter 与它覆盖的输入区域做逐元素乘法,再把这些乘积累加起来,得到输出 feature map 对应位置的一个值。这个过程捕捉了相邻特征之间的局部空间模式。

输入 feature map 通常有多个通道(channel)。以 RGB 图像为例,卷积不再只对 2D 空间区域求加权和,而是对 2D 空间区域 × 通道深度 一起求加权和。下图是 32×32 的三通道输入被 5×5×3 的 filter 卷成 28×28×1 输出的过程:



Source: CS231N https://cs231n.stanford.edu/slides/2025/lecture_5.pdf

每个 filter 产生一个输出通道。 要得到多个输出通道,就施加多个 filter。此外每个 filter 还带一个标量 bias,加到每个加权和上。

卷积算子(暂时忽略 bias)的输入输出可以总结为:

约束简化(与 Trainium2 相同):只支持 stride = 1;不需要处理 padding——输入 feature map 在送入你的 kernel 之前已经被 pad 好。

Mapping Convolution to Matrix Multiplication

课堂上讲过一种把"多 filter 卷积"变成一个大矩阵乘的做法:为每个空间 patch 单独造一行,也就是通常所说的 im2col(又叫 explicit GEMM)。

img

img2col

explicit GEMM :把 patch 物化成矩阵

slide 28(上面第二张图)画的就是 lowering(展开)这一步:

  • 展开矩阵:(H_out × W_out) × (K_h × K_w × C_in)
    • 行 = 一个输出位置;
    • 列 = 「patch 内的位置 × 输入通道」:按通道分成 C_in 组,每组 K_h·K_w 列(3×3 就是 9 列);
    • 越界的位置填 0 —— 图上边界那几行的 0 0 0 0 x00 x01 0 x10 x11 就是这种情况;
    • 图上最后一行 x00 x01 x02 x10 x11 x12 x20 x21 x22 是内部位置,9 个位置都是真实像素。
  • 权重矩阵:C_out × (K_h × K_w × C_in) —— 把 (C_out, C_in, K_h, K_w) 按「输出通道 × patch 位置 × 输入通道」展平(slide 右侧的 w_{filter}{位置}{通道})。
  • 一次矩阵乘就是整个卷积:
C_out × (H_out·W_out)  =  [ C_out × (K_h·K_w·C_in) ]  ×  [ (K_h·K_w·C_in) × (H_out·W_out) ]
     输出                        权重矩阵                        展开矩阵

也就是说,收缩维 K = K_h·K_w·C_in(3×3、C_in=256 时是 2304):窗口的 9 个位置被"折"进 K 维里,整个卷积是一次大矩阵乘,kernel 里没有任何 (i,j) 循环。

explicit GEMM
优点 一次大 GEMM,K 很长(矩阵乘单元最喜欢长 K);没有窗口循环
代价 展开矩阵是输入特征图的 K_h·K_w 倍(3×3 就是 9 倍):9 倍显存 + 9 倍读写;而且相邻输出位置的 patch 高度重叠(同一个输入像素最多出现在 9 个 patch 里),这份冗余是纯浪费

lecture 讲它是为了说明"卷积可以变成 GEMM",而不是说"必须先把那个矩阵造出来"。"物化版 im2col"在带宽受限的加速器上通常并不划算。

本作业为什么不用 explicit GEMM

本作业采用的是另一种 reduction;两者其实是"同一件事的两种展开方式"。K 维只取 C_in,把窗口的 (i,j) 放到外面当循环,每次对"错开 (i,j) 的输入"做一次 matmul,结果累加。Trainium2 版 README 对此有一条明确的注释,昇腾侧完全适用:

[!NOTE] This is a different conv -> matmul reduction than the one described in lecture that creates a separate row for every spatial patch.

Conv2D 与矩阵乘的数学形式对照:

本作业采用的 reduction 是:

  1. 把输入 feature map 的高、宽两个维度压平成一个维度,reshape 成 (Height × Width) × Input Channels;
  2. 然后用这个重塑后的输入依次乘以 filter 的每一个位置 (i, j)形成的一个矩阵,也就是图里的右边矩阵。注意这是最内层循环,固定 (i, j)(i 取 0 … Filter Height - 1,j 取 0 … Filter Width - 1);
  3. 每个 filter 切片的形状是 Input Channels × Output Channels,矩阵乘沿 Input Channels 维收缩。
  4. 为了把输入与当前 filter 切片对齐,输入必须按 (i, j) 做位移。
  5. 所有这些矩阵乘的结果累加起来,就是卷积的输出。

下面是该算法的伪代码(Trainium2 版 README 原文):

- Have the input image with shape (Input Channels, Image Height * Image Width)
- Have the filter weights with shape (Filter Height, Filter Weight, Input Channels, Output Channels)
- Initialize the output to appropriate shape of (Output Channels, Output Height * Output Width)

# Iterate over the filter height
for i in range(Filter_Height):
    # Iterate over the filter width
    for j in range(Filter_Width):

        # Shift the Input tensor by (i, j) to align with the filter's current position
        input_shifted = shift(input, (i, j))

        # Perform matrix multiplication between the input and the filter slice
        # Note that this is a full matmul, without limit on input sizes
        output += matmul(transpose(weight[i,j,:,:]), input_shifted)

[!NOTE] This is just an algorithmic description, and the purpose of this assignment is for you to figure out to map this algorithmic description to an efficient implementation on this hardware!

两种做法的关系:

explicit GEMM 本作业的 reduction
收缩维 K K_h·K_w·C_in(一次大 matmul) C_in(K_h·K_w 次 matmul 累加)
窗口位置 (i,j) 折进 K 维 外层循环
patch 矩阵 若物化则 9 倍显存 不物化

关键:两者都不需要真的把 patch 矩阵写到显存里,只要硬件/编译器能在"取数"这一步把窗口展开:

  • Trainium:shift(X, (i,j)) 是对 SBUF 中张量的切片(带偏移的读),不产生新数据;
  • 昇腾:有类似的API Load3D(文档原话 "image to column")在 L1→L0 的搬运路上把窗口展开进 L0A/L0B,同样不落显存。

所以"im2col"这个词要分开看:物化版 im2col(9 倍数据)谁都不划算;随路展开的 im2col(0 额外数据)正是昇腾 Load3D 做的事。

把这条算法落到昇腾上,有三处需要你自己解决。 Trainium2 版把这些留给了学生,但昇腾上的难点与 Trainium2 不同:

  1. shift(input, (i, j)) 怎么落到昇腾的取数指令上。 输入在显存里是 (C_in, H, W);按 (i, j) 位移后的二维块 (C_in, H_out × W_out) 是行不连续的——每读 W_out 个元素就要跳过 K_w − 1 个输入像素。

    昇腾上"把窗口展开成矩阵"有专门的硬件指令,但它与高阶 Matmul API 不可兼得:

    • 路线 A(直接对应):用 Load3D。 Load3D(即 LoadData 的 Load3Dv1/v2 形态,文档原话 "image to column")就是在 L1→L0 的搬运路上把特征图按窗口展开成二维矩阵写进 L0A/L0B:

      • filterW/filterH 给窗口大小

      • strideW/strideH 给滑动步长

      • padList 给补零

      • 窗口偏移由 kStartPt/mStartPt(V2)或 fetchFilterW/H、leftTopW/H(V1)给出。

 **位移只是一个参数——不需要在显存里预排,也不受 32 B 对齐限制**,这正是 `shift(X, (i,j))` 在昇腾上的对应物。

 代价:`Load3D` 写的是 **L0**,而 L0 归高阶 `Matmul` API 掌管 ⇒ **选它就必须降到低阶接口**(自己发 `LoadData` + `Mmad`,自己管 L1 分块、双缓冲与流水同步)。
  • 路线 B(留在高阶 API):把位移预先物化。 高阶接口的 A/B 只能来自显存、且必须是连续 ND 块(框架用 Nd2NzParams 随路转分形),所以每个列位移 j 都要在显存里各存一份"错列"数据(K_w 倍显存,3×3 即 3 倍;行位移 i 折进行号,不额外占空间)。此时 nValue / dValue / srcDValue 该怎么填,需要你自己推导并验证。
  1. 权重怎么放。 weight[i, j, :, :] 的 (C_out, C_in) 切片在两个平台上都不是连续内存(i、j 是最内层两个维度)。在 host 侧预先把权重重排是允许的(Trainium2 版 README 明确说明可以自由 reshape 输入)。两种路线各有自己的目标格式:

    • 路线 A(Load3D):重排成分形 FRACTAL_Z / NZ(常见形态 (C1_out, C1_in, K_h, K_w, C0_out, C0_in)),再由 Load3D 搬进 L0B;
    • 路线 B(高阶 Matmul):重排成 (K_h*K_w, C_out, C_in),让每次取 (i,j) 切片都是连续块(官方 conv 的 tiling 里 khL1/kwL1/nBL1 这些参数管的就是这件事)。

    另外注意:Trainium2 版保证"权重的总大小总能完整放进 SBUF(28 MiB)",这条保证在昇腾上不成立——fp16 下 9 个切片共 1.125 MiB,而 L1 只有 512 KB。单个切片(128 KB)放得下、全量放不下,循环顺序必须据此设计。

  2. K 循环放在哪一层。 这条算法的收缩维是 Input Channels。把它放进矩阵乘的 K 循环、让结果一直累加在 L0C 中直到该输出 tile 完全算完,是本作业推荐的循环顺序——这样每个输出 tile 只从 L0C 取出一次(路线 B 上是 Iterate(enPartialSum=true);路线 A 上是 Mmad 的累加模式)。

两条路线的取舍一览:

路线 A:Load3D + 低阶 路线 B:高阶 Matmul
shift(X,(i,j)) 怎么来 取数指令的窗口参数(元素级寻址) 显存里预先各存一份错列数据
额外显存 0 K_w 倍(3×3 → 3 倍)
tile 形状 2D 方块(halo 摊薄到 1.1~1.3×) 一整行(行 halo 3× 冗余)
权重要求 分形 FRACTAL_Z / NZ (K_h*K_w, C_out, C_in) 连续 ND
要自己写的 L1 分块、双缓冲、Mmad、流水同步 前置重排算子 + 位置函数
与 MatmulCallBackFunc 的关系 不需要 也救不了:回调只给"从哪个地址取"的自由,不给"格式"自由

Max Pool Layer Overview

数学定义与 Trainium2 相同。max pooling 在 feature map 上空间滑动一个 filter,但不计算加权和,而是取每个区域内的最大值。它逐通道独立进行,因此通道数不变。以 4×4 的三通道输入、2×2 的 pooling filter 为例,输出是 2×2×3:空间尺寸减半,通道数不变。

pool_size 的语义与 Trainium2 完全一致:一个参数同时表示 filter 大小与 stride,取值只能是 1 或 2。取 2 时如上图;取 1 时 maxpool 退化为恒等操作,输出与输入相同。取 1 看似无用,但它为融合层提供了灵活性——见 What You Need To Do。

昇腾侧的归约接口对应关系:

Trainium2 昇腾
nisa.tensor_reduce(nl.max, ...) ReduceMax / WholeReduceMax;或在 UB 上对配对元素用逐元素 Max
nisa.tensor_tensor(加 bias) Add / Adds(广播)
nisa.dma_copy DataCopy(尾块用 DataCopyPad)

Fusing Convolution and Max Pool

你要实现的是一个把卷积与 max pool 合并成单个融合操作的 Ascend C kernel。融合的目标与收益与 Trainium2 相同(省掉一次 25 MB 量级的中间结果往返)。昇腾上需要额外注意的是 AIC ↔ AIV 的协作方式:Cube 在 AIC、Vector 在 AIV,是两个核。

Cube 的结果怎么交给 Vector。 官方《融合算子编程》对 CV 融合算子的数据流定义如下:

融合算子数据流指融合算子的输入输出在各存储位置间的流向。以一个典型的Cube和Vector融合算子为例,逻辑位置间的数据流向如下图所示:

  • Cube的输出可以作为Vector的输入:CO2->VECIN
  • Vector的输出可以作为Cube的输入:VECOUT->A1->A2、VECOUT->B1->B2

其编程范式为五步:

  1. 初始化一个Matmul对象。
  2. 进行Matmul内部的计算。
  3. 将Matmul的计算结果搬运到Vector核上。
  4. 进行Vector矢量计算。
  5. 将输出结果搬运到Global Memory上。

用 Matmul 高阶 API 的默认 MIX 模式即可:GetTensorC(ubLocalTensor) 把 C 矩阵取到 UB,这一步由框架完成,不需要自己写 GM 中转或跨核同步。官方样例 matmul_leakyrelu 就是这个写法。

可以参考更多的融合算子案例:Matmul API中数据来源为VECOUT的矩阵乘实现方式

框架内部用的交换内存也在片上——《TSCM 输入的矩阵乘》写明"TSCM表示L1 Buffer空间对应的逻辑内存"。

融合与不融合的差别,官方也是这样对照的(同上《融合算子编程》):

  • 独立的矢量算子和矩阵算子实现:矩阵计算后的结果需要搬运到Global Memory上,然后由Global Memory搬运到Local Memory,再进行矢量算子的计算,计算和搬运都是串行执行,另外多个算子的调度执行,会增加算子的调度耗时。

  • 融合算子的实现方法:可以对数据进行切片,再通过流水的设计,使得矢量计算单元和矩阵计算单元实现并行计算;另外相比于不融合的单算子,减少了算子的调度耗时。

两种例外——只有这两种才需要自己处理跨核同步与 GM 中转,本作业不用:

写法 官方原文 结果去向
基础 API 手写 Mmad(分离模式) 《分离模式》:分离模式下,矩阵乘的计算结果从CO1(L0C Buffer)可以通过Fixpipe通路直接写入GM GM
Matmul 高阶 API + 双主模式 《AIC和AIV独立运行机制》:获取矩阵计算结果只支持调用IterateAll接口输出到GlobalTensor,即计算结果放置于Global Memory的地址,不能调用GetTensorC等接口获取结果 GM

下面这张图给出一个具体例子:6×6 输入、单输入通道,先做 stride = 1 的标准卷积,再对卷积结果做 2×2 的 max pool。

你的融合 kernel 接受以下参数(与 Trainium2 版一致):

  • X —— 一批输入图像,形状 (Batch Size, Input Channels, Input Height, Input Width)。保证 Input Channels 是 16 的倍数(昇腾改了口径,见下面说明)。
  • W —— 卷积 filter 权重,形状 (Output Channels, Input Channels, Filter Height, Filter Width)。保证 Filter Height == Filter Width,且 Output Channels 是 16 的倍数。

为什么昇腾上不是 128 的倍数?

Trainium 的 "128 的倍数" 来自 NKI 矩阵乘的 tile 形状上限:TILE_K = nl.tile_size.pmax = 128(partition 维,即收缩维 / 输入通道)与 TILE_M = nl.tile_size.gemm_stationary_fmax = 128(stationary 维,即输出通道方向)。形状必须被 128 整除,才能无 mask、无尾块地把矩阵切成硬件要求的固定尺寸块。

昇腾的 Cube 没有这种"形状上限",取而代之的是分形粒度 16×16×16(cube_m/n/k_size = 16)与 L0A / L0B / L0C 的容量上限:

约束来源 Trainium2 昇腾 A2 / A3
收缩维 / 输入通道(K 方向) ≤ 128(pmax),须被 128 整除 按你选的 baseK 分块,尾块由 Matmul API 内部补零累加,任意大小可算;硬性建议只是 16 的倍数(分形粒度;且 (C_in, N) 位移切片→ND2NZ 分形搬入时 nValue 推导才干净)
输出通道(M 方向) ≤ 128(stationary fmax),须被 128 整除 baseM 由 host tiling 任选,约束只有 L0C 容量(baseM × baseN ≤ 65536 fp16 / ≤ 32768 fp32);OC 余数交给尾块机制,同样不需要 128 倍数
官方佐证 — 默认 Norm 模板支持尾块;"无尾块"反而是 BasicBlock / SpecialBasicBlock 模板的特化前提(MatmulConfig 文档原文)

所以"OC / C_in 取 128 的倍数"在昇腾上只是一个可选调优(让 baseM = 128 / baseK = 128 恰好整除、零尾块浪费),不是硬性保证。本作业测试用例取 C_in = C_out = 256(= 2×128),两种口径下都是整除的。

⚠️ Trainium2 版额外保证"权重的总大小总能完整放进 SBUF",昇腾上没有这条保证(L1 只有 512 KB,装不下 fp16 的 1.125 MiB 全量权重)。这是昇腾版必须自己处理的差异。

  • bias —— 卷积 filter 的偏置,形状 (Output Channels)。
  • pool_size —— maxpool 的 filter 大小与 pooling stride。保证 (Input Height - Filter Height + 1) % Pool Size == 0。取 1 时融合 kernel 退化为普通卷积 kernel——这给了你"要不要池化"的灵活性。

输入输出的形状(忽略 bias):

out_height      = input_height - filter_height + 1
out_width       = input_width  - filter_width  + 1
out_pool_height = out_height // pool_size
out_pool_width  = out_width  // pool_size

X_out.shape == (Batch Size, Output Channels, out_pool_height, out_pool_width)

你可以自由 reshape 输入以换取更好的访问模式(Part 1 的 vector_add_stream 就是这么做的)。part2/conv2d_numpy.py 里给出了卷积层与 maxpool 层的 NumPy 参考实现,它描述了每一层的编程逻辑轮廓——把它融合成单个 kernel 正是本作业要你做的事。性能相关的提示见 General Tips 与 Tips on Using Ascend C。

What You Need To Do

本部分只聚焦一个文件:part2/conv2d.py(昇腾侧对应 part2/conv2d 算子工程)。里面给了基础骨架,你的任务是在 fused_conv2d_maxpool 中补完这个(融合的)Conv2D kernel。

General Tips

  • 正确性优先。 从最简单的用例开始:小图、无 bias、无 maxpool。kernel 在小图上正确之后,再让它支持"整张图放不下片上 Buffer"的大图;接着加入 bias 加法;最后把 max pool 融进 kernel。完全正确之后再开始做性能优化 / 加分项。
    • 测试框架按从易到难的顺序跑用例;你也可以先用不带 maxpool 的用例把卷积部分调到有性能,再回来处理融合的 max pool。
  • 把算法吃透。 在考虑任何 tiling 策略之前,先确认你真正理解了上面那条算法需要的三种矩阵操作(乘法、位移、累加)。然后把矩阵与它们的维度画出来,思考如何映射到硬件上——尤其是内存层次。
    • 你可能需要预处理输入数组(reshape 或 transpose)以获得更好的访问模式。
    • 提示:想想昇腾的矩阵乘接口相比 NKI 多要求了什么——L1 → L0A/L0B 的显式 staging,以及 L0A/L0B 的 512 B 对齐分形格式。
  • 时刻盯住 tile 尺寸。 你不可能一次算完整张输出,所以必须决定把输出的哪个维度切开。回忆片上约束:UB 的 32 B 对齐、L0A/L0B 的 512 B 对齐、L0C 的容量上限。定下输出 tile 形状之后,反推你需要 X 和 W 的哪个子集。
  • 排循环顺序时要考虑数据局部性。 你需要的 for 循环来自多个源头:算法本身的 filter 高与宽、分块矩阵乘、以及 batch。
    • 排好之后,推荐让中间结果留在 L0C 里,直到该 tile 的计算完全结束——这样输出数组的每一部分只被写一次,输出局部性更好(其他做法也可能达到相近性能)。
    • 在此之上,再排列剩余循环以优化输入局部性。拿不准就试几种数据访问模式,并想想为什么。
  • 用 profiler 指导调优。 有了能跑的 kernel 之后,大概率还需要继续调才能拿满/加分。msprof 是你的朋友:找 Tensor/Cube 空闲、利用率低的大段空隙,重构代码把这些空隙消掉。
    • 回想 Part 1 里你优化向量加和转置 kernel 的经验,思路是通的。

Testing

用提供的测试框架验证你的实现。进入 part2/ 目录执行:

python3 test_harness.py

要检查融合了 maxpool 的 Conv2D kernel 的正确性与性能,加上 --test_maxpool 参数。

测试框架先跑正确性测试、再跑性能检查。注意:Trainium2 给出的参考性能阈值(4596 μs / 1002 μs)是特定硬件的实测值,在昇腾上无意义,必须自建基准,方法如下:

  • 基准 A(推荐):CANN 内置算子 aclnnConvolution + aclnnMaxPool 组合的耗时(aclnn 系列接口无需 IR 定义即可直调);
  • 基准 B:你自己先写一版朴素实现(im2col + Matmul 高阶 API)的耗时,作为对照;
  • 评分线:达到基准 A 的 1.2 倍以内 = 满分,与原作业的 120% 语义一致。

测试框架会用 float32 与 float16 两种数据类型调用你的 kernel,且 fp16 的性能要求更严。

单精度与半精度的差异在昇腾上除了算力,还会影响 L0C 能容纳的 tile 大小(fp32 下 baseM × baseN 的上限减半)——设计 tiling 时要把两种精度都考虑进去。

Writeup and Profiling

需要提交一份报告,简要描述你的实现,并说明你是如何做优化的。报告里必须给出实测的 MFU(Model FLOPs Utilization),float16 与 float32 两种数据类型都要。在昇腾上用 msprof 采集:

msprof op --soc-version=<你的SoC> ...

[!TIP] 两套剖析工具(Trainium 的 neuron-profile 与昇腾的 msprof)口径不同,不能逐条直接比较。报告里请显式给出Trainium 指标 → 昇腾指标的对照表,并注明机型(npu-smi 的 Chip Name + NPU Name)、CANN 版本与 --soc-version。

MFU 用昇腾自己的口径算:

FLOPs = 2 × B × C_out × C_in × K_h × K_w × H_out × W_out
MFU   = FLOPs / (实测耗时 × 该型号的峰值算力)

本作业性能用例(B=1, C_in=C_out=256, 224×224, k=3,故 H_out=W_out=222)代入:

FLOPs      = 2 × 1 × 256 × 256 × 9 × 222 × 222 ≈ 5.81e10 = 58.1 GFLOP
最小访存量 ≈ 输入 25.7 MB + 权重 1.18 MB + 输出 25.2 MB ≈ 52 MB
算术强度   ≈ 58.1e9 / 52e6 ≈ 1117 FLOP/Byte   →  严重 compute-bound

峰值算力必须从实例的实际规格读(存储容量用 GetCoreMemSize(),算力看实例规格页或 npu-smi info),不要照抄 Trainium2 的 79 TFLOPS。作为目标参考:Trainium2 的 fp16 满分线 1002 μs 反推约 58 TFLOPS ≈ 79 TFLOPS 峰值的 73%,这个量级可以当作昇腾调优的目标。

结论:Part 2 的优化重心是让 Cube 少空转(tiling + 流水 + AIC/AIV 重叠),而不是省带宽——这与 Part 1 的向量加/转置(访存 bound)正好相反,是个很好的对比。

Tips on Using Ascend C

  • 对齐是第一优先级:UB 32 B、L0A/L0B 512 B、L0C/BT/FP 64 B。W_out = H_out = 222 不是 16 的倍数(222 = 13×16 + 14),尾块必须用 DataCopyPad 或 padding 处理——这是昇腾版比 Trainium2 版多出来的第一道坎。
  • 不要把配置值硬编码:UB / L1 / L0C 的容量在运行时用 GetCoreMemSize() 取,核数用 GetCoreNumAic() / GetCoreNumAiv() 取,blockDim 不要写死。
  • SPMD 要显式切分:用 GetBlockIdx() / GetBlockNum() 决定当前核处理哪一块。
  • 注意 UB 的 bank 冲突:多个 Tensor 同时活跃时,按 256 B 粒度对齐分配(Part 0)。
  • 先在 CPU 域跑通(ICPU_RUN_KF 孪生调试),再上 NPU 实例测性能——这是省卡时的关键,也是昇腾相对 Trainium2 版多出来的一条工程路径。
  • 优先复用官方样例:融合算子的起点是 matmul_leakyrelu;矩阵编程的完整样例集在 cann/asc-devkit 的 examples/01_simd_cpp_api/ 下。

Extra Credit

除 Part 1 Step 3 的额外加分外,Ascend 侧的加分空间集中在 Part 2 的融合深度与性能:

  • 把 bias 与 maxpool 真正融进流水:不要"算完整个卷积再回来做池化",而是让 AIV 处理上一块的同时 AIC 已经在算下一块(生产者-消费者重叠);
  • 对齐的代价:W_out = H_out = 222 不是 16 的倍数,尾块处理得好不好会直接反映在利用率上;
  • 小图优化:32×16 这类小图上,核数与 tile 的选择空间与 224×224 完全不同——两种尺寸都跑出好数字才算调透。

Grading Guidelines

评分点的分布与 Trainium2 保持一致(Part 1 共 30 分、Part 2 共 70 分),但性能类评分必须重新标定基线:

  • 正确性部分原样保留:小图 / 大图 / 大图+bias / 大图+bias+maxpool 四类用例,参考实现逐元素 np.allclose 对拍,正确性是拿分的前提;
  • 性能部分改为相对基准:以 aclnnConvolution + aclnnMaxPool 的耗时为基准 A,达到基准 A 的 1.2 倍以内 = 满分(与原作业的 120% 语义一致);
  • 两种精度都要测:float32 与 float16,且 fp16 的要求更严;
  • 报告里必须给出你选的 baseM/baseN/baseK 及其依据、shift 的落地方式、以及实测 MFU。

Hand-in Instructions

与 Trainium2 相同:提交代码与报告。报告中的每个性能数字都必须给出真实的执行命令与输出,并注明机型(npu-smi 的 Chip Name + NPU Name)、CANN 版本与 --soc-version。


参考资料

融合算子与 Cube ↔ Vector 交接

这几篇是判断"Cube 的结果要不要落 HBM"的依据。结论:用 Matmul 高阶 API + 默认 MIX 模式时,Cube 结果经 GetTensorC 直接进 UB(逻辑位置 CO2 → VECIN),不需要落 HBM。

  • 融合算子编程 —— 核心篇:CV 融合算子的五步范式,以及 Cube的输出可以作为Vector的输入:CO2->VECIN 这条逻辑通路;文中明确把"融合算子"与"结果先搬 GM 再搬回的独立算子"对立起来讲
  • 分离模式 —— 用基础 API 手写 Mmad 时耦合 / 分离模式的差异;分离模式下 L0C 的出口只有 Fixpipe → GM,耦合模式才能走 UB
  • AIC 和 AIV 独立运行机制 —— 双主模式(enableMixDualMaster=true)下只能 IterateAll 输出到 GM、不能 GetTensorC,是为提性能而牺牲片上交接的一档
  • Matmul 使用说明 —— MIX 模式(默认)、A/B/C/Bias 四者的 TPosition 配置、REGIST_MATMUL_OBJ
  • IterateAll 接口 · GetTensorC 接口 —— 取 C 矩阵的两条接口;ubCmatrix 形参支持 TPosition::VECIN / TSCM
  • TPosition —— CO1 / CO2 / VECIN / TSCM 等逻辑位置的含义与映射
  • Fixpipe 接口 —— 随路类型转换 / 量化反量化 / ReLU / NZ→ND;表里写明 A2/A3 仅支持 FixpipeParamsV220(通路 CO1 → GM)
  • TSCM 简介 —— 判断"要不要经过 HBM"的关键一篇:Cube 与 Vector 之间通过队列通信;mask 位 2 说明"VECTOR0 与 VECTOR1 映射到 CUBE 上的内存位置"
  • TSCM 输入的矩阵乘 —— 一句话定性:"TSCM 表示 L1 Buffer 空间对应的逻辑内存",即框架的交换内存是 AIC 侧片上 L1,不是 HBM
  • GetLibApiWorkSpaceSize 接口 —— Matmul 需要的 workspace 是固定的库级尺寸,不随矩阵规模变化
  • CrossCoreSetFlag · CrossCoreWaitFlag —— 分离模式核间同步(模式 2 = AIC↔AIV);仅基础 API 手写路径或手动做 AIC/AIV 流水重叠时需要,融合算子用高阶 API 时由框架处理
  • GetCurNpuArch 接口 —— 查硬件平台的 __NPU_ARCH__ 版本号(100x / 200x / 220x / 300x / 351x),决定哪些接口可用

融合算子官方样例(asc-devkit / samples)

  • matmul_leakyrelu —— CV 融合算子的标准写法:__global__ __mix__(1, 2)、C 矩阵 TPosition::VECIN、Iterate<true>() → GetTensorC<true>(ubLocalTensor, ...) → Vector 就地计算,全程无手写 GM 中转
  • matmul_mixdualmaster —— 双主模式样例(性能取向,代价是结果只能落 GM)
  • MmadInvocation —— 基础 API 手写矩阵乘,可用来看分离模式与耦合模式的具体差异

results matching ""

    No results matching ""