diff --git a/docs/lab/Lab3.5-AscendC-Op/image/ai-core-basic.png b/docs/lab/Lab3.5-AscendC-Op/image/ai-core-basic.png new file mode 100644 index 00000000..68c41c84 Binary files /dev/null and b/docs/lab/Lab3.5-AscendC-Op/image/ai-core-basic.png differ diff --git a/docs/lab/Lab3.5-AscendC-Op/image/ai-core-coupled.png b/docs/lab/Lab3.5-AscendC-Op/image/ai-core-coupled.png new file mode 100644 index 00000000..29b1ce10 Binary files /dev/null and b/docs/lab/Lab3.5-AscendC-Op/image/ai-core-coupled.png differ diff --git a/docs/lab/Lab3.5-AscendC-Op/image/ai-core-separated.png b/docs/lab/Lab3.5-AscendC-Op/image/ai-core-separated.png new file mode 100644 index 00000000..d581b763 Binary files /dev/null and b/docs/lab/Lab3.5-AscendC-Op/image/ai-core-separated.png differ diff --git a/docs/lab/Lab3.5-AscendC-Op/image/ascendc-copyin-size.png b/docs/lab/Lab3.5-AscendC-Op/image/ascendc-copyin-size.png new file mode 100644 index 00000000..c09d3e79 Binary files /dev/null and b/docs/lab/Lab3.5-AscendC-Op/image/ascendc-copyin-size.png differ diff --git a/docs/lab/Lab3.5-AscendC-Op/image/ascendc-copyout-size.png b/docs/lab/Lab3.5-AscendC-Op/image/ascendc-copyout-size.png new file mode 100644 index 00000000..4b4cd5d4 Binary files /dev/null and b/docs/lab/Lab3.5-AscendC-Op/image/ascendc-copyout-size.png differ diff --git a/docs/lab/Lab3.5-AscendC-Op/image/ascendc-double-buffer.png b/docs/lab/Lab3.5-AscendC-Op/image/ascendc-double-buffer.png new file mode 100644 index 00000000..10543bee Binary files /dev/null and b/docs/lab/Lab3.5-AscendC-Op/image/ascendc-double-buffer.png differ diff --git a/docs/lab/Lab3.5-AscendC-Op/image/ascendc-pipeline.png b/docs/lab/Lab3.5-AscendC-Op/image/ascendc-pipeline.png new file mode 100644 index 00000000..94d86941 Binary files /dev/null and b/docs/lab/Lab3.5-AscendC-Op/image/ascendc-pipeline.png differ diff --git a/docs/lab/Lab3.5-AscendC-Op/image/ascendc-ub-reuse.png b/docs/lab/Lab3.5-AscendC-Op/image/ascendc-ub-reuse.png new file mode 100644 index 00000000..52e31ac4 Binary files /dev/null and b/docs/lab/Lab3.5-AscendC-Op/image/ascendc-ub-reuse.png differ diff --git a/docs/lab/Lab3.5-AscendC-Op/image/empty_kernel_overhead_vs_aiv_cores.svg b/docs/lab/Lab3.5-AscendC-Op/image/empty_kernel_overhead_vs_aiv_cores.svg new file mode 100644 index 00000000..e4450f89 --- /dev/null +++ b/docs/lab/Lab3.5-AscendC-Op/image/empty_kernel_overhead_vs_aiv_cores.svg @@ -0,0 +1 @@ +00.511.522.5Task Duration (µs)0510152025303540启动 AIV 核数0.67 µs2.25 µs \ No newline at end of file diff --git a/docs/lab/Lab3.5-AscendC-Op/image/msprofop.png b/docs/lab/Lab3.5-AscendC-Op/image/msprofop.png new file mode 100644 index 00000000..f33f47f4 Binary files /dev/null and b/docs/lab/Lab3.5-AscendC-Op/image/msprofop.png differ diff --git a/docs/lab/Lab3.5-AscendC-Op/image/msprofopsim-left.png b/docs/lab/Lab3.5-AscendC-Op/image/msprofopsim-left.png new file mode 100644 index 00000000..8de58332 Binary files /dev/null and b/docs/lab/Lab3.5-AscendC-Op/image/msprofopsim-left.png differ diff --git a/docs/lab/Lab3.5-AscendC-Op/image/msprofopsim-right.png b/docs/lab/Lab3.5-AscendC-Op/image/msprofopsim-right.png new file mode 100644 index 00000000..6724a31b Binary files /dev/null and b/docs/lab/Lab3.5-AscendC-Op/image/msprofopsim-right.png differ diff --git a/docs/lab/Lab3.5-AscendC-Op/image/score.png b/docs/lab/Lab3.5-AscendC-Op/image/score.png new file mode 100644 index 00000000..c3bf9d05 Binary files /dev/null and b/docs/lab/Lab3.5-AscendC-Op/image/score.png differ diff --git a/docs/lab/Lab3.5-AscendC-Op/index.md b/docs/lab/Lab3.5-AscendC-Op/index.md new file mode 100644 index 00000000..1960ffd0 --- /dev/null +++ b/docs/lab/Lab3.5-AscendC-Op/index.md @@ -0,0 +1,680 @@ +# Lab 3.5:昇腾算子开发与优化 + +!!! info "实验信息" + + 负责助教:刘天洋,徐晨,周冠一,胡笠桁,胡哲文 + + !!! tip "Bonus 实验" + + 这是一个 **Bonus 实验**,不强制要求同学们完成。欢迎学有余力、对国产计算硬件感兴趣的同学挑战。 + + !!! note "硬件资源赞助" + + 本实验的昇腾 NPU 计算资源由 **华为技术有限公司** 赞助。在此致以诚挚谢意。 + +## 实验目的 + +在其他实验中,我们了解了 x86-64、ARM 以及 RISC-V 等架构的 CPU 和 NVIDIA GPU 的硬件架构,并尝试写出更充分利用硬件特性的程序以优化运行效率。本次实验将把视野扩展到国产计算生态的代表硬件 **昇腾(Ascend)NPU**:你需要在昇腾 910B4 NPU 上实现一个简单的 `fused_add_rmsnorm` 算子,在保证功能正确的前提下,尽可能地提高性能。 + +完成本实验后,你应当能够: + +- 了解昇腾 910 架构 NPU 的基础架构与计算单元组织; +- 了解国产计算生态基本情况,至少掌握 AscendC / TileLang / Triton 中一种程序的基本开发能力; +- 能够比较昇腾 NPU 与 NVIDIA GPU 在设计理念上的异同; +- 掌握华为昇腾 NPU 上常见的算子优化手段,并理解这些手段背后对应着硬件的哪些特性; +- 使用 `msprof` 等工具对算子进行性能分析,并据此定位瓶颈、指导优化方向。 + +!!! note "前置课程" + + 7 月 14 日「华为 CANN」课程覆盖了本实验的相关内容。 + + 如果你对相关知识点不熟悉,也可以回顾课程内容。或参考文末的参考资料。 + +## 背景知识 + +### FusedAddRmsNorm + +`FusedAddRmsNorm`,顾名思义,是把「残差加法(Add)」和「RMSNorm」融合在一起的一个算子。 + +先来看 [RMSNorm](https://arxiv.org/abs/1910.07467),这是一种对向量做均方根归一化的方法。对一个长度为 $H$ 的向量 $x$,它先用均方根做缩放,再乘以可学习权重 $w$: + +$$ +\operatorname{RMSNorm}(x) = \frac{x}{\sqrt{\dfrac{1}{H}\sum_{i=1}^{H}x_i^2 + \varepsilon}} \odot w, +$$ + +其中 $w$ 是与 $x$ 逐元素相乘的可学习缩放权重,$\varepsilon$ 是为数值稳定加入的小常数。RMSNorm 直接用均方根做缩放,计算较轻,在 Transformer 等模型中被广泛用作归一化层。 + +而在 Transformer 前向中,RMSNorm 之前往往紧跟一个残差加法(residual add),把上一路的输入累加到残差流上。这其实就是一个矢量加法。`fused_add_rmsnorm` 就是把这两步**融合成一个算子**,避免残差结果在全局显存中多写一次再读回。融合后的计算为: + +$$ +\begin{aligned} +R &= x + \mathit{residual}, \\ +\mathit{rms} &= \sqrt{\dfrac{1}{H}\sum_{i=1}^{H}R_i^2 + \varepsilon}, \\ +y &= \dfrac{R}{\mathit{rms}} \odot w. +\end{aligned} +$$ + +即该算子同时输出两部分: + +- `residual_out` $= R = x + \mathit{residual}$,供后续残差流继续使用; +- `y` $= \operatorname{RMSNorm}(R)$,作为本层的归一化输出。 + +### 昇腾 910 NPU + +#### 昇腾 NPU 与达芬奇架构 + +昇腾(Ascend)NPU 是华为自研的 AI 加速器,其核心 IP 是 **达芬奇(Da Vinci)架构**。达芬奇架构采用 **Cube + Vector + Scalar** 三种计算单元的异构组合: + +- **Cube 单元**:承担矩阵乘法(GEMM)类运算,一个周期可完成一次 $M\times K\times N$ 的矩阵乘(具体维度由代际决定),类似 CPU 上的 AMX / SME 扩展; +- **Vector 单元**:承担向量与元素级运算(Cast、Add、Mul、Reduce、Sqrt 等),类似 CPU 上的 AVX / SVE 向量化扩展。但宽度较宽(一次处理 256B); +- **Scalar 单元**:承担标量运算、地址计算与控制流。 + +Cube + Vector + Scalar 三个物理单元在一个 **AI Core** 内部并行执行,配合多级片上存储(UB / L1 / L0A / L0B / L0C)完成计算。多个 AI Core 组成一颗 NPU,通过 HBM 全局显存共享数据。 + +#### AI Core 与 AIC / AIV 分离架构 + +在 910B 系列上(包括 A2 推理和训练系列),AI Core 采用了 **AIC / AIV 分离**的设计: + +- **AIC(AI Cube)**:一个 AIC 核内置 Cube 单元、L1、L0A/L0B/L0C 等存储,主要承担矩阵乘类指令; +- **AIV(AI Vector)**:一个 AIV 核内置 Vector 单元、UB(Unified Buffer)等存储,主要承担向量与元素级指令; +- AIC 与 AIV 在物理上分离,通过片内总线通信;它们之间以及 AI Core 之间通过全局显存(Global Memory, GM / HBM)协作。 + +
+
+ ![AI Core 分离架构](image/ai-core-separated.png) +
+
AI Core 分离架构示意
+
+ +每个核内还有多级片上存储:AIV 侧的 **UB(Unified Buffer,向量计算的快速访存区)**、AIC 侧的 L1、L0A/L0B/L0C 等。多个 AI Core 之间相互独立,通过 tiling 把数据切分到各核上并行执行。 + +!!! note "AI Core 耦合架构" + + 事实上,AI Core 也存在耦合架构,即可以直接在 AIV 的 UB 与 AIC 的 L1 Buffer / Buffer L0C 之间搬运数据而不需要经过全局内存。这种架构常见于边缘推理系列(如昇腾 310B 系列等)。 + +#### AIV 内部结构 + +!!! warning "本实验不涉及 AI Cube" + + 由于算子本身不涉及矩阵乘法相关的运算,同时 AIC 算子的编写门槛较高,我们不会对 AIC 相关的开发框架和硬件特性做展开,感兴趣的同学可以自行了解。 + +每个 AIV 核内关键资源如下(on Ascend 910B4 NPU): + +| 资源 | 容量/规格 | 用途 | +| --- | --- | --- | +| **UB(Unified Buffer)** | 192 KiB / AIV | 向量计算的快速访存区,所有 Vector 指令的操作数必须在 UB 上 | +| **Vector 单元** | 256B(VLEN) | 执行 Cast/Add/Mul/Reduce 等向量指令 | +| **MTE2 搬运单元** | — | GM → UB 的数据搬入 | +| **MTE3 搬运单元** | — | UB → GM 的数据搬出 | +| **Scalar 单元(S)** | — | 标量运算、地址计算、控制流 | + +!!! note "数据搬运与流水" + + AIV 上有 MTE2(GM→UB)、V(Vector 计算)、MTE3(UB→GM)三条独立的流水线。在本次实验中,同学们需要了解以下两种同步机制: + + 1. **`TQue` 队列语义自动同步**:当你用 `inQueX.AllocTensor` → `DataCopy` → `inQueX.EnQue` 发射数据搬入 UB 指令后,再 `inQueX.DeQue` 取出来给 Vector 用时,**`EnQue` / `DeQue` 内部会自动插入 MTE2→V 的同步事件**——你不需要手动 `SetFlag/WaitFlag`。同理,输出端 `outQueY.EnQue` → `outQueY.DeQue` 也会自动处理 V→MTE3 的同步。这一机制主要保证内存搬运和向量计算流水线之间的依赖被正常处理。 + + 2. **`PipeBarrier` 显式屏障**:这一屏障用于保证某一条流水线内部的同步关系,常用的是 `PIPE_V`。Ascend NPU 的 V 流水线被设计为直接操作内存。Ascend NPU 不保证 V 流水线内部的 RAW 依赖能够被自动、正确地处理。开发者需要在存在 RAW 依赖的指令之间显式调用 PipeBarrier(),以确保本核此前发出的所有 V 指令均已执行完成。 + 为了减少频繁插入屏障带来的编程负担,bisheng 编译器提供了 --cce-auto-sync 选项。启用后,编译器会根据可见的 LocalTensor 读写依赖自动插入必要的屏障。不过,当代码中涉及指针运算、容器传递或手工地址操作时,编译器的依赖分析可能失效,此时就需要同学们手动插入屏障。 + + 更多有关同步的信息,感兴趣的同学可以参考 [Ascend C API列表](https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/latest/API/ascendcopapi/atlasascendc_api_07_0003.html). + +### 算子开发路径 + +本实验框架支持 Ascend C、TileLang 和 Triton 三种算子开发路径。总体而言,Ascend C 更接近昇腾硬件,开发者需要显式处理数据切分、片上存储、数据搬运和计算流水;TileLang 与 Triton 则提供了更高层的 DSL 编程模型,由编译器承担更多底层实现工作。 + +- **Ascend C** Ascend C 是昇腾原生的 C/C++ 算子编程语言与 API 体系。它通过 DataCopy、Cast、Add、Mul、BlockReduceSum 等接口屏蔽部分指令级细节,同时保留对核间切分、片上存储、数据搬运、队列、同步和流水的细粒度控制,适合需要深入调优硬件性能的场景。一个完整的 Ascend C 自定义算子通常由三部分组成: + + - **算子原型(Op 定义)**:声明算子的输入、输出与属性,注册推理 shape 与数据类型的推导函数,以及 tiling 函数。其中 OpDef 类与注册宏等骨架可由 CANN 的 `msopgen` 工具从算子信息 JSON 自动生成,但 tiling 函数中的切分逻辑仍需开发者手写; + - **Host 侧 tiling**:在 Host 上根据输入 shape、数据类型和平台信息,计算切分参数,决定每个核处理多少数据、UB 如何分配,并把 tiling 数据传给 kernel(类比于 CUDA 中设置 grid/block 维度并准备 kernel 参数的 host 侧代码,但额外承担了数据分区与 UB 规划); + - **Kernel 实现**:在 AI Core 上执行真正的计算,负责 GM↔UB 搬运、UB 内计算、结果写回(类比于 CUDA 中的 kernel 函数)。 + +- **TileLang** 在 Lab 3 中已登场,是一种面向高性能计算的 DSL。它把 tiling、shared memory、pipeline 等概念直接暴露为语言原语,同样可编译到昇腾 NPU 后端(见 [tilelang-ascend](https://github.com/tile-ai/tilelang-ascend))。相较 Ascend C,它省去了手写 host tiling 与 GM↔UB 搬运的样板代码,但在某些细节上对硬件的控制力可能弱于原生路径。 + +- **Triton** 是基于 block 的算子编程模型,通过 [triton-ascend](https://github.com/triton-ascend/triton-ascend) 后端可编译到昇腾 NPU。它以 `tl.program_id` 切分任务、以 block 为单位组织计算,写法接近你在 GPU 上的经验,便于把已有的 GPU kernel 快速迁移过来。 + +三种路径最终编译出的都是跑在 Ascend NPU 上的机器码,性能上限由硬件决定。路径之间的差异在于「你能在多大程度上接近硬件」。 + +!!! warning "多种实现路径" + + 我们提供了 AscendC、TileLang、Triton 三种实现路径的代码框架。你只需要选择一种路径实现即可,**多种实现不会带来进一步加分!** + + 同时,为了使大家接触和了解国产计算生态,我们**鼓励大家在本实验中使用 AscendC 完成你的算子**。代码框架中我们只针对 AscendC 路径提供基线实现,实验文档也会以 AscendC 算子开发为主。 + + 但是同样需要注意,使用 AscendC 路径**不会自动给你带来额外加分**。 + +## 性能分析 Profiling + +性能优化离不开测量。只有从 profiling 结果找到根据并针对性分析,才能找到更合理有效的优化方向。 + +对标 NVIDIA 的 Nsight System,Nsight Compute 等性能分析工具,华为也提供了 MindStudio 一族工具,命令行主要体现为 `msprof` 类指令,而桌面图形化版可以前往对应位置[下载 MindStudio Insight](https://www.hiascend.com/document/detail/zh/mindstudio/latest/GUI_baseddevelopmenttool/MindStudioInsight/docs/zh/user_guide/overview.md)。 + +本实验主要涉及以下两个 Profile 指令: + +### `msprof op`:分析算子的宏观特点 + +`msprof op` 对目标算子进行板端采集。主要的性能分析结果集中在宏观指标,包括算子运行时间,NPU 上各个单元(如 Vector 等)的利用率,UB 等不同层级缓存间的带宽,L2 Cache 的命中率,不同 AI Core 任务量是否均匀等。 + +!!! tip "一个简单的例子" + + 运行以下指令即可进行一次标准的算子性能数据采集: + + ```bash + # 采集默认的计算、访存、流水线和资源冲突指标 + msprof op \ + --kernel-name="fused_add_rms_norm" \ + --launch-count=1 \ + --aic-metrics=Default \ + --output=./op_prof \ + python3 checker/test_op.py 2 + ``` + + 将生成目录中的 `visualize_data.bin` 导入 MindStudio Insight 可以查看图形化结果。 + +分析结果时可以考虑检查以下问题: + +- `Block Dim` 与各核工作量是否符合 tiling 设计,是否存在空闲核、明显拖尾或负载不均匀; +- Vector 等运算单元的利用率是否较高; +- UB 高速缓存的利用带宽是否够高,次一级的 L2 Cache 命中率是否足够。 + +!!! note "如何读取分析的数据" + + 利用率表示该运算单元处于活跃状态的比例,不一定等同于『有效工作』的比例。多条流水线可以重叠,因此不同单元之间的利用率没有明显关系;单个指标较高或较低也不足以单独证明瓶颈。 + +!!! example "`msprof op` 结果可视化示例" + + 以下图片展示了一个 AIV 部分的 Profile 结果,可以看到不同层级缓存之间利用的带宽,Vector 单元的利用率,以及 L2 Cache 的命中率等。 + +
+ msprof op 结果可视化示例 +
`msprof op` 在 MindStudio Insight 中的可视化结果。
+
+ +### `msprof op simulator`:观察指令流水与同步 + +Simulator 在 CPU 上模拟指定 SoC 的指令执行,适合检查指令发射顺序、流水线空泡、`SetFlag/WaitFlag` 等同步事件和资源冲突。主要生成的结果是一个流水图(也叫 Timeline)。 + +它不是真实 NPU 计时工具:仿真耗时不能与 `Task Duration(us)` 比较,缓存、带宽竞争和运行时调度也应以板端结果为准。 + +!!! tip "一个简单的例子" + + 运行以下指令即可进行一次标准的算子模拟数据采集: + + ```bash + msprof op simulator \ + --soc-version=Ascend910B4 \ + --kernel-name="fused_add_rms_norm" \ + --launch-count=1 \ + --output=./op_sim \ + python3 checker/test_op.py 2 + ``` + + 由于是 CPU 侧模拟算子在 NPU 上的行为,所以我们需要显示指定 -soc-version 为 Ascend 910B4。 + +!!! warning "不要尝试模拟真实数据规模" + + 由于是 CPU 模拟,实际上 `msprof op simulator` 的性能非常糟糕,考虑到我们通过这个指令得到的结果主要是用来分析不同单元流水线之间的运转情况,模拟时请**设置较小的数据规模**,以便在合适的时间和资源使用下完成模拟。 + +!!! example "`msprof op simulator` 流水线示例" + + 以下两张流水图表明了同一个计算功能在不同实现下的模拟流水图,其中上三角形是各种同步 flag。你应该可以明显看到右图中的 Vector 利用率相比于左图有明显提升(少了很多空泡和停顿)。 + +
+
+ 优化前 msprof op simulator 流水图 +
优化前:Vector 流水存在较多空泡与同步停顿
+
+
+ 优化后 msprof op simulator 流水图 +
优化后:Vector 流水更连续,同步空泡明显减少
+
+
+ + 不同流水线之间的 Timeline 还可以展示是否产生了流水线重叠执行(即同一时间不同硬件单元同时工作),详见思考题。 + +### 从现象提出可验证的假设 + +Profiler 给出的是证据,而不是自动生成的优化结论。可以按“现象—假设—实验”的方式推进: + +| 观测现象 | 可优先检查 | 下一步实验 | +| --- | --- | --- | +| MTE2/MTE3 时间较长,Vector 经常等待 | 搬运粒度过小、重复读写 GM、未形成流水 | 合并搬运或复用 UB 数据,通过流水开启双缓冲技术 | +| Vector 流水长时间连续工作 | 指令数量、规约层数、类型转换或高代价指令 | 优化计算逻辑,实现更高性能的同样计算功能等 | +| 各流水线都存在大段空泡 | 过宽的屏障、Scalar 依赖、tile 太少 | 用 simulator 对齐空泡与同步事件,再收窄同步范围 | +| 部分核明显更早结束 | 多核切分不均或尾块集中 | 调整 Tiling 和 NPU 侧 Init 时的工作分配 | +| Roofline 接近带宽上限 | 算子可能受 GM 带宽约束 | 优先减少 GM 往返,而不是继续增加 Vector 指令并行度 | + +每次只改变一个主要因素,先跑完整正确性测试,再重复计时。只有当对应指标与端到端耗时同时朝预期方向变化时,才能较有把握地解释优化收益。 + +## Ascend C 简介 + +Ascend C 通过 C++ 模板与 intrinsic 暴露 AIV 上的存储、搬运与计算指令。下面只列出本实验会接触到的核心概念,作为阅读 baseline 代码和继续编写 Ascend C 代码的参考。 + +### Tensor 编程模型 + +在 Ascend C 的编程模型中,主要操作的数据被视作一个 Tensor。 + +!!! tip "你也可以操作单个数据" + + Ascend C 也提供了单个数据的 API(如 `GetValue` 和 `SetValue`)。但这两个 API 的性能通常比较糟糕。 + +- **UB 与 LocalTensor**。Vector 指令的操作数必须位于 AIV 内的 UB(Unified Buffer)上。Ascend C 以 `LocalTensor` 表示一段 UB 上的逻辑数据视图,承载实际数据类型 `T`(如 `half`、`float`)与长度信息,本身不直接分配内存,需从 `TBuf` 或 `TQue` 中获取。`GlobalTensor` 则是 GM 上的对应视图,用于 `DataCopy` 等搬运 API 的源/目的。 + +- **TPosition**。`TPosition` 是一个枚举,标识缓冲区在存储层级中的物理位置,常见的三种取值如下: + + - `VECIN`:与 MTE2 流水衔接的输入侧 UB,通常作为搬入队列位置; + - `VECOUT`:与 MTE3 流水衔接的输出侧 UB,通常作为搬出队列位置; + - `VECCALC`:纯 Vector 侧的 UB,不参与 `TQue` 的跨流水线同步,常用于常驻或中转数据。 + +- **TBuf**。`TBuf` 是 UB 上的静态缓冲区,由 `TPipe::InitBuffer` 在 kernel 启动前按字节大小分配,生命周期与 kernel 一致。它直接持有 `LocalTensor`,不参与任何同步事件的插入,适合存放常驻数据(如 weight)或临时中转缓冲(如平方缓冲、规约暂存)。 + +- **TQue 与流水**。`TQue` 在 `TBuf` 之上提供队列语义,用 `AllocTensor → EnQue → DeQue → FreeTensor` 描述一段 LocalTensor 在流水阶段间的所有权与同步事件转移。`BUFFER_NUM` 控制队列深度,设为 2 即为相邻 tile 的搬入与计算提供交叠空间。 + +### 内存搬运类 API + +`DataCopy` 用于 GM↔UB 间对齐、连续或规则分段的数据搬运; + +`DataCopyPad` 进一步支持非 32B 对齐的尾块,可通过 `isPad`、`leftPadding`/`rightPadding`、`paddingValue` 控制尾部的填充值,避免越界读写或破坏对齐。`DataCopyExtParams` 以 `blockCount` / `blockLen` / `srcStride` / `dstStride` 描述多段搬运的几何。 + +### Vector 计算 API + +Vector 侧的元素级与逐元素运算包括 `Cast`(类型转换,可指定 `RoundMode`)、`Add`、`Mul`、`Div`、`Sqrt`、`Rsqrt` 等,通常以 `dst, src0, src1, maskCount` 形式调用;`Duplicate` 用标量广播填充一段 LocalTensor。所有调用都按 32B 对齐(FP16 16 元素、FP32 8 元素)操作,非对齐尾数由 mask 或 `DataCopyPad` 处理。 + +沿一维规约的指令以 `Reduce` 为主。可以指定规约长度等参数,规约操作包括求和,取极值等。同时 AscendC 提供了更底层的 `BlockReduce` 与 `WholeReduce` 操作,用来减小 `Reduce` 的开销。感兴趣的同学可以进一步了解这些 API 的使用方法。 + +!!! tip "查阅官方 API 参考" + + 上述 API 的支持数据类型、对齐、mask、repeat、临时空间与同步的要求,以及功能和使用方法的详细描述,都在 Ascend C 官方 API 文档里有所说明,实际编写算子时,请及时查阅相关 API 文档。详见 [Ascend C API 列表](https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/latest/API/ascendcopapi/atlasascendc_api_07_0003.html)。 + + +## 编写高性能的 Ascend C 算子 + +!!! warning "不要照抄优化思路!" + + 以下给出的是可能有效的优化方向,而不是必须实现的。请始终以性能结果和 profiling 结果为准。 + +### 让搬运与计算形成流水 + +AIV 的 MTE2、Vector 和 MTE3 可以分别执行搬入、计算和搬出。由于这三个单元本身是相互独立的,因此如果在同一时间重叠不同单元的执行将能够大幅掩盖时间,获得加速。Ascend C 的编程模型中倡导把一个 tile 的处理拆成 `CopyIn → Compute → CopyOut` 三段式,明确生命生产者和消费者关系,并用 `TQue` 在阶段间传递 LocalTensor 的所有权与同步事件。这样的编写范式能够更好的让编译器识别依赖关系,并且**自动**开启流水。 + +
+
+ ![Ascend C CopyIn、Compute、CopyOut 三段流水编程范式](image/ascendc-pipeline.png) +
+
Ascend C 三段流水编程范式。图源:昇腾社区
+
+ +将队列的 buffer 数设为 2(或者更多),可以为相邻 tile 交叠提供空间:当一份数据参与 Vector 计算时,另一份数据才有机会同时搬入或搬出。这个也被称作 Double Buffering 技术。 + +!!! note "不一定更快" + + Double Buffer 并不总是保证自动加速,同时其是否生效还需要满足相邻阶段之间没有不必要的全局同步等条件。 + +### 让中间结果留在 UB + +融合算子的主要价值,是让前一阶段的输出直接成为后一阶段的 UB 输入,而不是每经过一步就写回 GM。官方示例中的反例执行两次 GM 往返,正例则在 `VECCALC` 中保留中间值。 + +
+
+ ![连续 Vector 计算中经由 GM 中转与在 UB 内复用中间结果的对比](image/ascendc-ub-reuse.png) +
+
连续 Vector 计算的中间结果应尽量在 UB 内衔接。图源:昇腾社区
+
+ +对本算子而言,$R=x+\mathit{residual}$ 随后还要用于输出 `residual_out`、平方和规约以及计算 $y$。只要 UB 容量允许,就应保留可复用的 $R$,并把写回 `residual_out` 的副本交给输出队列,避免为后续计算再次读取 GM。`weight` 在所有行上相同,也可以按核或按 tile 复用。与此同时,应统计所有常驻 buffer 与双缓冲队列的总占用,避免为了复用而挤压 tile 大小或导致 UB 超限。 + +### 根据 shape 设计 Tiling + +多核和单核 tiling 需要一起设计: + +- 运行时通过平台接口查询 AIV 核数与 UB 容量,不把某个环境中的数值写死到策略中; +- 各行相互独立时,优先沿独立维(如 $B$)切分。使用商和余数分配大小核,使各核行数最多相差 1,避免简单向上取整造成尾部空闲核; +- 当独立维很小时,盲目启动全部 AIV 没有收益。沿其他维(如 $H$)切分虽然能增加并行度,却会引入跨核规约或额外 pass,需要根据实测决定是否值得; +- tile 大小既要满足 UB 预算,也要让每核有足够多的 tile 支撑流水。最大的 tile 不一定最快,最小的 tile 也会增加循环和搬运启动开销。 + +??? example "从启动开销的角度思考为什么不一定需要启动全部 AIV" + + 在 NPU 中,每个核在被启动时都要单独付一次初始化开销(加载配置、建立执行上下文、准备片上资源等),且这部分几乎无法通过流水、搬运或规约优化消除——它只和「启动了多少个核」有关。下面是用一个 body 几乎为空的算子,只改变 `<<>>` 下发 AIV 数目,用 `msprof op` 测得的 Task Duration,可以近似认为是启动开销: + +
+ ![空 kernel 的运行时间和启用 AIV 数目的关系](image/empty_kernel_overhead_vs_aiv_cores.svg) +
空 kernel 的运行时间和启用 AIV 数目的关系
+
+ + 它随启动核数近似线性增长(约 2 核 0.67 µs → 40 核 2.25 µs),本身处在微秒量级,在大算子里几乎被 profiling 波动淹没;但当算子整体只有几微秒、单核计算量又很小时就占比可观。因此计算量过小的任务要选择「够用就好」的核数,而不是无脑写满全部 AIV。 + +公开评测 shape 中既有 $B=1$,也有 $H$ 非 32B 对齐的情况。一个只对 $[256,1024]$ 快、但在小 $B$ 或尾块上出错的 tiling 不是有效实现。 + +### 提高搬运效率 + +MTE 搬运存在固定启动成本,单次搬运过小时难以充分利用带宽。 + +??? example "单次搬运量与有效带宽的关系" + + 以下是昇腾社区对于不同搬运量带宽利用率的一个测试,可以看到小的数据搬运并不能很好的利用带宽。相关结果数据并不一定适用于我们的 910B4 NPU 上,只是表明一个相关趋势。 + +
+
+ HBM 到 UB 的单核带宽随单次搬运数据量变化 +
HBM → UB
+
+
+ UB 到 HBM 的单核带宽随单次搬运数据量变化 +
UB → HBM
+
+
+ +

图源:昇腾社区

+ +若 UB 预算和数据布局允许,可以尝试用一次 `DataCopy` 搬运多行,或用 `blockCount`、`blockLen` 与 stride 参数表达规则的多段搬运,减少逐行发射指令的开销。与此同时需要注意: + +- GM 起始地址、每行 stride 和 UB 地址的对齐会共同影响搬运效率; +- 非对齐尾块应使用 `DataCopyPad` 或等价 mask 正确处理,不能为了对齐越界读写; +- 合并搬运会增加 UB 占用,并可能减少每核 tile 数,应与 Double Buffer 一起重新评估; + +### 选择合适的规约实现与精度 + +高效的归约操作是本实验中这个算子能够获得高性能的关键一环。Ascend C API 中提供的 `Reduce` 类算子可能存在较多的同步 Flag 设置和边界检测情况,实际上可能并不能获得很好的性能。因此,Ascend C API 同时提供了更加底层的 `BlockReduce` 操作和 `WholeReduce` 操作。两者在单次执行速度和吞吐效率上各有优劣,但是合作起来就可能获得比原生 `Reduce` 更好的规约性能和运算单元利用率。 + +同时,本实验的输入输出是 FP16,但这并不意味着平方和也适合在 FP16 中累加。使用 FP32 作为中间结果通常更稳健,但会占用更多 UB,也可能增加 Cast 和 Vector 计算开销。可以尝试减少重复 Cast、缩短 FP32 数据的存活范围; + +!!! danger "禁止一味地使用低精度来获得更高性能" + + 任何低精度或近似方案都必须通过我们最后的精度校验! + +??? tip "拓展阅读" + + - [Ascend C 算子性能优化实用技巧 01:流水优化](https://www.hiascend.com/zh/developer/techArticles/20240819-1) + - [Ascend C 算子性能优化实用技巧 02:内存优化](https://www.hiascend.com/zh/developer/techArticles/20240823-1) + - [Ascend C 算子性能优化实用技巧 03:搬运优化](https://www.hiascend.com/zh/developer/techArticles/20240906-1) + - [Ascend C 算子性能优化实用技巧 04:Tiling 优化](https://www.hiascend.com/zh/developer/techArticles/20240920-1) + - [Ascend C 算子性能优化实用技巧 05:API 使用优化](https://www.hiascend.com/developer/techArticles/20241107-1) + +## 实验任务 + +你的任务是在昇腾 910B4 NPU 上实现并优化 `FusedAddRmsNorm` 算子。你可以从 **Ascend C、Triton-Ascend、TileLang-Ascend** 中任选一种开发路径;只需完成一种实现,尝试多种路径不会带来额外加分。 + +实现需要同时满足以下要求: + +- 保持给定的输入、输出和属性接口; +- 通过公开与隐藏测试的正确性检查; +- 在指定评测配置下尽可能优化 kernel 性能; + +### 接口与计算语义 + +设 $B$ 为行数,$H$ 为隐藏层宽度: + +| 张量 | 参数类型 | 形状 | 数据类型 | 含义 | +| --- | --- | --- | --- | --- | +| `x` | 输入 | $[B,H]$ | FP16 | 输入 | +| `residual` | 输入 | $[B,H]$ | FP16 | 残差 | +| `weight` | 输入| $[H]$ | FP16 | RMSNorm 缩放权重 | +| `eps` | 输入| 标量 | FP16 | 防止除零的微小偏移 | +| `y` | 输出|$[B,H]$ | FP16 | 归一化结果 | +| `residual_out` |输出| $[B,H]$ | FP16 | 残差加法结果 | + +对每一行 $b\in[0,B)$,算子计算: + +$$ +\begin{aligned} +R_b &= x_b + \mathit{residual}_b, \\ +\mathit{rms}_b &= \sqrt{\dfrac{1}{H}\sum_{i=0}^{H-1}R_{b,i}^{\,2} + \varepsilon}, \\ +y_b &= \dfrac{R_b}{\mathit{rms}_b} \odot w, \\ +\mathit{residual\_out}_b &= R_b. +\end{aligned} +$$ + +三个开发路径对 checker 暴露的入口均为 `fused_add_rmsnorm(x, residual, weight, eps)`,其中 `eps` 默认为 $10^{-6}$。 + +??? note "关于 `enable_pdl` 属性" + + Ascend C 的底层算子定义中还保留了一个布尔属性 `enable_pdl`,用于与 FlashInfer 原始签名保持一致。FlashInfer 中该属性控制是否启用 NVIDIA CUDA 12.3+ 引入的 **Programmatic Dependent Launch(PDL,程序化依赖启动)**——允许后一个 kernel 在前一个 kernel 完全结束前就开始发射与初始化,以掩盖 kernel launch 延迟。 + + 在昇腾 NPU 上不存在对等的硬件机制,因此本实验中该属性不产生任何实际效果:PyTorch 扩展胶水层会固定向算子传入 `false`,Triton 与 TileLang 路径也无需暴露该属性。你不需要在实现中处理它,也不应期望通过它获得性能收益。 + +!!! warning "本实验采用非原地接口" + + FlashInfer 原始接口会原地修改输入,而本实验规定 `x` 和 `residual` 为只读输入,并返回独立的 `y` 与 `residual_out`。请勿改变这一约定。 + +### 修改范围与限制 + +你可以修改所选路径目录下的实现,并在该目录内增加必要的辅助文件: + +| 开发路径 | 可修改目录 | +| --- | --- | +| Ascend C | `src/ascendc/` | +| Triton-Ascend | `src/triton/` | +| TileLang-Ascend | `src/tilelang/` | + +允许自行设计 kernel、tiling、核数、片上存储布局和针对不同 shape 的实现分支。禁止: + +- 修改 `checker/`、`env.sh`、输入生成、golden 或计时逻辑; +- 调用已有 RMSNorm、FusedAddRmsNorm 或等价高层算子代替被测计算; +- 硬编码测试数据、隐藏 shape 或输出结果; +- 利用评测程序漏洞绕过计算或正确性检查; +- 依赖课程环境中未提供的额外软件包或自建工具链。 + +可以参考开源实现,但需要在报告中注明来源,并说明自己的实现与修改。 + +## 代码框架 + +实验代码位于仓库的 `src/lab3p5/`: + +```text +src/lab3p5/ +├── env.sh # 加载课程 CANN 与 Python 环境 +├── README.md # 代码框架使用与提交说明 +├── checker/ +│ ├── build.sh # 构建并安装 Ascend C 算子 +│ ├── run.sh # 正确性检查 +│ ├── profile.sh # 固定性能 case 的 msprof 采集 +│ ├── test_op.py # 输入生成、FP32 golden 与逐元素比较 +│ ├── case_specs.py # 公开测试配置 +│ └── get_time.py # 解析 op_summary 中的 kernel 时间 +└── src/ + ├── __init__.py + ├── ascendc/ + │ ├── op_host/ + │ │ ├── CMakeLists.txt + │ │ └── fused_add_rms_norm.cpp # 算子注册与 Host tiling + │ ├── op_kernel/ + │ │ ├── CMakeLists.txt + │ │ ├── fused_add_rms_norm.cpp # Device kernel + │ │ └── fused_add_rms_norm_tiling.h # Tiling 数据结构 + │ ├── extension/custom_op.cpp # PyTorch 扩展胶水 + │ ├── common/pytorch_npu_helper.hpp # PyTorch NPU 辅助头文件 + │ ├── CMakeLists.txt / CMakePresets.json + │ ├── build_op.sh # 构建 Ascend C 算子与 wheel + │ └── setup.py + ├── triton/ + │ ├── __init__.py + │ └── fused_add_rmsnorm.py # Triton kernel 与 launcher + └── tilelang/ + ├── __init__.py + └── fused_add_rmsnorm.py # TileLang kernel 与 launcher +``` + +### 选择开发路径 + +| 路径 | 建议先阅读 | 运行方式 | +| --- | --- | --- | +| Ascend C | `src/ascendc/op_host/fused_add_rms_norm.cpp`、`src/ascendc/op_kernel/fused_add_rms_norm.cpp` | 直接运行 checker;脚本会在当前任务中按需构建并安装算子 | +| Triton-Ascend | `src/triton/fused_add_rmsnorm.py` | 提交任务时设置环境变量 `LANG=triton` | +| TileLang-Ascend | `src/tilelang/fused_add_rmsnorm.py` | 提交任务时设置环境变量 `LANG=tilelang` | + +Ascend C 路径提供了一个以正确性为主的 baseline:Host 侧读取 shape 和属性并生成 tiling;Device 侧按行分配工作,FP16 数据进入 UB 后主要使用 FP32 计算,再转换为 FP16 输出。它已经展示了 `TQue`、`TBuf`、数据搬运、Vector 计算和规约的基本组织方式,但同步、流水和 tiling 仍有优化空间。 + +本文不逐项复述 Ascend C API。使用接口前,请直接查阅对应 CANN 版本的 [Ascend C API 参考](https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/latest/API/ascendcopapi/atlasascendc_api_07_0003.html),尤其注意支持的数据类型、对齐、mask、repeat、临时空间和同步要求。 + +### 构建与自测 + +所有命令都应在 `src/lab3p5/` 下执行。NPU 任务通过 `hpc submit` 提交到 `lab3p5` 分区;`run.sh` 和 `profile.sh` 会自行加载 `env.sh`,一般不需要在提交任务前手动 `source`。 + +```bash +# 正确性:运行全部公开 case +hpc submit -p lab3p5 bash checker/run.sh +hpc submit -p lab3p5 -e LANG=triton bash checker/run.sh +hpc submit -p lab3p5 -e LANG=tilelang bash checker/run.sh + +# 正确性:只运行一个公开 case;编号从 1 开始 +hpc submit -p lab3p5 bash checker/run.sh 2 + +# 性能:固定 shape 性能测试(256×1024),不接受 case 参数 +hpc submit -p lab3p5 bash checker/profile.sh +hpc submit -p lab3p5 -e LANG=triton bash checker/profile.sh +hpc submit -p lab3p5 -e LANG=tilelang bash checker/profile.sh +``` + +`checker/run.sh` **只负责正确性检查**;`checker/profile.sh` **只采集 student 算子的性能**,使用 `msprof op --warm-up=10` 并输出一次 `Task Duration(us)`。进行性能测试前,应先用 `run.sh` 验证正确性。 + +Ascend C 路径下,每次改动代码后需要自己进行编译,编译脚本已经写好为 `checker/build.sh`,在 Devpod 内即可进行。TileLang 和 Triton 路径会自动触发编译,不需要同学们手动进行。 + +## 如何获取计算资源 + +我们通过[实验平台](https://platform.s.zjusct.io)提供 **arm64-910b** DevPod,容器拉取的镜像中已经配有本实验所需的 CANN 工具链和 Triton/Tilelang 包环境,一般不需要自行安装工具链。更详细的平台和文件同步说明见[集群使用](https://hpc101.zjusct.io/guide/)。你需要做的包括: + +1. 登录实验平台; +2. 创建预设为 `arm-910b` 的 DevPod; +3. 在 DevPod 中获取课程仓库并进入 `src/lab3p5/`; +4. 执行 `source ./env.sh` 后开始构建和测试。 +5. 在需要在 NPU 上执行算子时使用 `hpc submit -p lab3p5 ` 即可提交至分配有一张 910B4 NPU 的计算分区执行你的命令。 + +!!! warning "家目录不共享!" + + 华为为我们提供的 Ascend 910B4 8 卡裸金属机器地理上分布于华北-乌兰察布地区。距离超算队在杭州的集群和其他硬件资源有一定的地理距离。由于 NFS 对于时延有较高要求,本平台的家目录不会与其他硬件资源的家目录共享。 + +!!! danger "请注意区分 Devpod" + + 在创建 Devpod 时,Lab 4 任务一所需的鲲鹏环境对应的 Devpod 预设为 `arm64-920b`,而 Lab3.5 的预设 Devpod 为 `arm64-910b`,两者环境和家目录均不互通,**请注意区分**。 + +## 评分方式 + +评测包括**正确性**和**性能**两部分。只有通过正确性检查的实现才会进入性能计分。 + +### 正确性验证 + +我们提供多个 Case 测试你的算子。公开 case 由 `checker/case_specs.py` 定义。下表同时给出代码中的 case 索引和命令行编号: + +| case 索引 | 命令行编号 | $B\times H$ | 说明 | +| --- | --- | --- | --- | +| 0 | 1 | $32\times4096$ | 小规模、对齐 | +| 1 | 2 | $256\times1024$ | 性能评测配置 | +| 2 | 3 | $1\times4096$ | 单行 | +| 3 | 4 | $1997\times3037$ | 行数与尾部均不对齐 | +| 4 | 5 | $2048\times4096$ | 大规模、对齐 | + +输入由固定 seed 在运行时**随机**生成。正式评测还会使用不同 shape 的隐藏 case(保证数据范围大致一致),因此不能只针对公开配置硬编码实现。 + +Golden 在 FP32 下完成残差加法、平方和、均值、开方、除法和权重缩放,最后转换为 FP16。`y` 与 `residual_out` 均需逐元素通过检查。 + +对参考值 $g_i$ 和输出值 $o_i$,元素在满足下列任一条件时通过(相对精度或绝对精度小于给定阈值): + +$$ +|o_i-g_i|\le 10^{-3} +\quad\text{或}\quad +\frac{|o_i-g_i|}{\max(|g_i|,10^{-12})}\le 10^{-3}. +$$ + +整个张量要求错误元素比例为 0。一个更快但未通过全部正确性检查的实现不会获得性能分数。 + +### 性能评分 + +为了方便 OJ 等实现,性能测试只涉及单个 Shape,评测配置为: + +- **Shape**:$[256,1024]$; +- **输入输出类型**:FP16; +- **`eps`**:$10^{-6}$; +- **指标**:被测 kernel 的 `Task Duration(us)`(来自 `msprof op` 的结果,热身 10 次)。 + +基于我们的基线和优化结果,目前的评分曲线如下(对数曲线): + +
+ 性能评分曲线 +
性能评分曲线(横轴为 kernel 耗时,纵轴为得分)。
+
+ +我们提供的 Ascend C 基线性能为起始评分点,满分 120 分,超出的 20 分将作为 Bonus。 + +## 实验报告要求 + +实验报告提交 PDF,重点说明你如何从测量得到优化决策,不需要重复大段背景知识或 API 文档。报告至少应包含: + +- 使用的开发路径、测试环境、软件版本和运行命令; +- 算子的计算过程、数据依赖和初始实现; +- baseline 的正确性、性能数据和 profiling 证据; +- 每项主要优化针对的瓶颈、关键修改及其收益; +- 最终正确性结果和性能结果; +- 尝试过但未采用的方案及原因; +- [思考题](#思考题)作答; +- 参考过的资料或开源实现。 + +较长的代码、完整 profiler 输出和构建日志不必全部放入正文,只需保留能够支持结论的部分。 + +!!! tip "失败尝试也值得记录" + + 没有带来加速的尝试可以帮助说明原先的瓶颈判断、优化的副作用,或不同指标之间的取舍。 + +!!! danger "关于 AI 使用" + + 本实验允许使用 AI Agent 辅助开发和理解资料,但最终报告应由你自行组织和核实,禁止使用 AI 生成。 + +## 思考题 + +!!! tip "这里可能没有标准答案" + + 部分思考题可能没有一个标准的正确答案。我们更希望看到你在实验过程中遇到的实际情况以及你个人的思考和理解,即使可能存在错误。 + +1. 对 $[256,1024]$、FP16 配置,估算算子必须进行的 GM 读写量和主要浮点操作数。说明你的计数口径,计算算术强度,并结合 `msprof` 结果判断实现更接近计算瓶颈还是访存瓶颈。 +2. 你的算子是否开启了 Double Buffer 流水?请你通过 `msprof op simulator` 的结果向我们证明。如果你开启了双缓冲流水,请说明你是显示编写了依赖还是让编译器自动识别开启的? +3. (Bonus)我们讲到了 `TQue` 相关的概念,结合你对相关文献的阅读和实际的编程以及 Profiling 体验,`TQue` 本身是否是一个真实存在的队列?我们在 `EnQue` 和 `DeQue` 时相关的 LocalTensor 是否发生了在队列之间的拷贝或移动? +5. (Bonus)今年第一季度,华为正式推出了新一代 NPU Ascend 950PR,引入了很多新的变化。请阅读相关资料,向我们展示一下相比 910 系列,950PR 带来了哪些新的硬件特性?你也可以比较昇腾 NPU(910B 系列)与 NVIDIA GPU 在硬件设计理念,开发语言等方面的异同。或者搜索一下一代代昇腾 NPU 和 NVIDIA GPU 的迭代过程,你有什么发现? +6. (Bonus)今年是我们首次将 NVIDIA GPU 以外的其他异构计算平台引入到课程和实验中,欢迎你给我们分享你的体验,对于昇腾 NPU 的使用体验或者相关开发语言(Ascend C、Triton-ascend、TileLang-ascend)的使用体验等(锐评也可以,不会扣分的!)。 + +## 提交要求 + +需要提交: + +1. **实现代码**:所选开发路径对应的完整目录; +2. **实验报告**:单独的 PDF 文件。 + +请在课程平台对应的 Lab 3.5 任务中分别上传代码目录和报告 PDF;具体文件名与上传入口以任务开放时的页面说明为准。 + +代码上传规则如下: + +| 开发路径 | 上传目录 | 平台放置位置 | +| --- | --- | --- | +| Ascend C | `src/ascendc/` | `src/ascendc/` | +| Triton-Ascend | `src/triton/` | `src/triton/` | +| TileLang-Ascend | `src/tilelang/` | `src/tilelang/` | + +新增辅助文件必须位于所选路径目录内,并保证代码能在课程提供的干净环境中构建和运行。 + +请勿提交: + +- `checker/`、`env.sh` 或对这些文件的修改; +- `build_out/`、`dist/`、wheel、custom OPP 安装目录等构建产物; +- `PROF*`、`op_prof/`、`op_sim/`、`prof_out/` 等 profiling 原始目录; +- 测试数据、缓存或其他与实现无关的大文件。 + +## 参考资料 + +### 算法与参考实现 + +- [Root Mean Square Layer Normalization](https://arxiv.org/abs/1910.07467) +- [FlashInfer: Kernel Library for LLM Serving](https://github.com/flashinfer-ai/flashinfer) +- [FlashInfer `fused_add_rmsnorm` 文档](https://docs.flashinfer.ai/generated/flashinfer.norm.fused_add_rmsnorm.html) + +### Ascend C 与性能分析 + +- [Ascend C 开发文档](https://www.hiascend.com/document/detail/zh/canncommercial/80RC31alpha001/devguide/ascendc) +- [Ascend C API 参考](https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/latest/API/ascendcopapi/atlasascendc_api_07_0003.html) +- [Ascend C 最佳实践](https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/850alpha002/opdevg/ascendcbestP/atlas_ascendc_best_practices_10_0031.html) +- [msProf 算子调优工具](https://www.hiascend.com/document/detail/en/canncommercial/850/devaids/optool/atlasopdev_16_0082.html) +- [CANN 软件下载与文档](https://www.hiascend.com/software/cann) + +### 其他开发路径 + +- [Triton 文档](https://triton-lang.org/) +- [triton-ascend](https://github.com/triton-ascend/triton-ascend) +- [TileLang 文档](https://tilelang.com/) +- [tilelang-ascend](https://github.com/tile-ai/tilelang-ascend) diff --git a/src/lab3p5/.gitignore b/src/lab3p5/.gitignore new file mode 100644 index 00000000..202a7200 --- /dev/null +++ b/src/lab3p5/.gitignore @@ -0,0 +1,4 @@ +build*/ +*egg*/ +.whl +__pycache__/ diff --git a/src/lab3p5/README.md b/src/lab3p5/README.md new file mode 100644 index 00000000..04e5b696 --- /dev/null +++ b/src/lab3p5/README.md @@ -0,0 +1,56 @@ +# FusedAddRmsNorm(Lab 3.5) + +请选择 Ascend C、Triton-Ascend 或 TileLang-Ascend 中的一种实现 `fused_add_rmsnorm`。正式评测会使用隐藏 case,请勿针对公开 shape 硬编码。 + +## 提交 NPU 任务 + +在 `lab3p5/` 根目录使用 `hpc submit` 将任务提交到 `lab3p5` NPU 分区。任务默认使用当前目录作为工作目录。 + +### Ascend C + +```bash +hpc submit -p lab3p5 bash checker/run.sh # 正确性测试 +hpc submit -p lab3p5 bash checker/profile.sh # 性能测试 +``` + +主要修改 `src/ascendc/op_kernel/` 和 `src/ascendc/op_host/` 下的算子代码。 + +### Triton-Ascend + +```bash +hpc submit -p lab3p5 -e LANG=triton bash checker/run.sh +hpc submit -p lab3p5 -e LANG=triton bash checker/profile.sh +``` + +实现文件:`src/triton/fused_add_rmsnorm.py`。 + +### TileLang-Ascend + +```bash +hpc submit -p lab3p5 -e LANG=tilelang bash checker/run.sh +hpc submit -p lab3p5 -e LANG=tilelang bash checker/profile.sh +``` + +实现文件:`src/tilelang/fused_add_rmsnorm.py`。 + +如需后台提交,可添加 `-d`,再使用 `hpc logs -f ` 查看日志。 + +## Checker 说明 + +- `checker/run.sh` 只执行正确性测试;不传参数时运行全部公开 case。 +- `hpc submit -p lab3p5 bash checker/run.sh 2` 可单独运行 case 2,即 `256×1024`。 +- `checker/profile.sh` 只测试并输出 student 算子的性能。 +- 性能测试固定使用 case 2(`256×1024`),不接受 case 参数。 +- 性能采集使用 `msprof op --warm-up=10`,最终输出一次 `Task Duration(us)`。 + +## 提交代码 + +验证完成后,通过 HPC101 平台上传所选语言对应的整个目录: + +| 实现 | 提交目录 | +| --- | --- | +| Ascend C | `src/ascendc/` | +| Triton-Ascend | `src/triton/` | +| TileLang-Ascend | `src/tilelang/` | + +只提交其中一个实现目录。不要提交 `checker/`、`env.sh`、`README.md`、构建产物或 profiling 输出目录。 diff --git a/src/lab3p5/checker/build.sh b/src/lab3p5/checker/build.sh new file mode 100644 index 00000000..887553a8 --- /dev/null +++ b/src/lab3p5/checker/build.sh @@ -0,0 +1,58 @@ +#!/bin/bash +# Build + install the FusedAddRmsNorm Ascend C op AND the pybind extension. +# Run AFTER `source env.sh` (from the kit root). Lives under `checker/` now; +# ROOT still resolves to the kit root (parent of this dir) so paths match the +# runtime layout the README documents. +# +# IMPORTANT: this script runs in a subshell, so the custom_opp set_env.bash it +# sources (registers the aclnn op with libopapi.so via ASCEND_CUSTOM_OPP_PATH +# / LD_LIBRARY_PATH) does NOT leak to the caller. After this returns, either +# `source ./env.sh` again, or `source $CUSTOM_OPP_HOME/vendors/customize/bin/ +# set_env.bash` to make aclnnFusedAddRmsNorm resolvable at runtime. run.sh does +# the former automatically. +set -e +ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)" +cd "$ROOT/checker" + +OP_DIR="$ROOT/src/ascendc" +export CUSTOM_OPP_HOME="${CUSTOM_OPP_HOME:-$HOME/custom_opp}" +# env.sh must be sourced first so the system CANN (ASCEND_TOOLKIT_HOME) is on +# PATH — we don't second-guess its location here. +: "${ASCEND_TOOLKIT_HOME:?env.sh not sourced — run 'source ./env.sh' first (ASCEND_TOOLKIT_HOME unset)}" + +echo "=== [build] patch op CMakePresets -> system CANN ===" +python3 - <", cv["ASCEND_CANN_PACKAGE_PATH"]["value"]) +PY + +echo "=== [build] build + install custom op ===" +cd "$OP_DIR" +rm -rf build_out +export ASCEND_HOME_PATH="${ASCEND_HOME_PATH:-$ASCEND_TOOLKIT_HOME}" +export DDK_PATH="$ASCEND_HOME_PATH" +export NPU_HOST_LIB="$ASCEND_HOME_PATH/$(arch)-$(uname -s | tr '[:upper:]' '[:lower:]')/devlib" +bash build_op.sh 2>&1 | tail -6 +RUN_PKG="$(ls build_out/custom_opp_*.run 2>/dev/null | head -1)" +[[ -z "$RUN_PKG" ]] && { echo "[build] no .run produced"; exit 1; } +"$RUN_PKG" --quiet --install-path="$CUSTOM_OPP_HOME" 2>&1 | tail -2 +# set_env.bash (generated above) already prepends op_api/lib to LD_LIBRARY_PATH +# and sets ASCEND_CUSTOM_OPP_PATH — no manual LD_LIBRARY_PATH poking needed. +# shellcheck disable=SC1091 +source "$CUSTOM_OPP_HOME/vendors/customize/bin/set_env.bash" +export CUSTOM_OPP_SOURCED=1 + +echo "=== [build] build + install pybind wheel ===" +cd "$OP_DIR" +# setup.py + extension/custom_op.cpp + common/pytorch_npu_helper.hpp all live +# under src/ascendc/ now (setup.py is a sibling of extension/ and common/). +python3 setup.py build bdist_wheel +pip3 install dist/custom_ops-*.whl --force-reinstall --no-deps + +echo "[build] done." diff --git a/src/lab3p5/checker/case_specs.py b/src/lab3p5/checker/case_specs.py new file mode 100644 index 00000000..7b1975d9 --- /dev/null +++ b/src/lab3p5/checker/case_specs.py @@ -0,0 +1,18 @@ +"""Public case specifications for FusedAddRmsNorm. + +Each case is just a spec — inputs are generated on the fly from `seed` (so the +kit ships NO .bin data), and the golden is computed in-process (fp32, then cast +fp16) by test_op.py. The judge swaps this file for a hidden case_specs.py with +different shapes; test_op.py itself never changes. + +Tuple layout: (B, H, eps, data_range, seed) + data_range in {"S": (-1,1), "M": (1,10), "L": (-1000,1000)} +""" + +CASES = [ + (32, 4096, 1e-6, "S", 1001), # small, aligned + (256, 1024, 1e-6, "S", 1002), # regular, aligned + (1, 4096, 1e-6, "S", 1003), # single row — scalar-bound + (1997, 3037, 1e-6, "S", 1004), # misaligned B + misaligned H (tail handling) + (2048, 4096, 1e-6, "S", 1005), # large, aligned +] diff --git a/src/lab3p5/checker/get_time.py b/src/lab3p5/checker/get_time.py new file mode 100644 index 00000000..a8d65cf8 --- /dev/null +++ b/src/lab3p5/checker/get_time.py @@ -0,0 +1,48 @@ +"""Read the single Task Duration(us) produced by ``msprof op``. + +``msprof op`` performs warm-up internally and writes one selected operator to +``OPPROF_*/OpBasicInfo.csv``. This parser deliberately rejects zero or multiple +rows instead of silently mixing different operators or shapes. + +Usage: + python3 checker/get_time.py +""" +import csv +import sys +from pathlib import Path + +def read_task_duration(root: Path) -> float: + csv_paths = sorted(root.rglob("OpBasicInfo.csv")) + if len(csv_paths) != 1: + raise RuntimeError( + f"expected one OpBasicInfo.csv under {root}, found {len(csv_paths)}" + ) + + with csv_paths[0].open("r", encoding="utf-8-sig", newline="") as f: + rows = list(csv.DictReader(f)) + durations = [ + row.get("Task Duration(us)", "").strip() + for row in rows + if row.get("Task Duration(us)", "").strip() + ] + if len(durations) != 1: + raise RuntimeError( + f"expected one Task Duration(us) in {csv_paths[0]}, found {len(durations)}" + ) + return float(durations[0]) + + +def main(): + if len(sys.argv) != 2: + print(f"Usage: {sys.argv[0]} ", file=sys.stderr) + sys.exit(2) + try: + duration = read_task_duration(Path(sys.argv[1])) + except (OSError, ValueError, RuntimeError) as exc: + print(f"[ERROR] {exc}", file=sys.stderr) + sys.exit(1) + print(f"{duration:.4f}") + + +if __name__ == "__main__": + main() diff --git a/src/lab3p5/checker/profile.sh b/src/lab3p5/checker/profile.sh new file mode 100755 index 00000000..c4e8572d --- /dev/null +++ b/src/lab3p5/checker/profile.sh @@ -0,0 +1,75 @@ +#!/bin/bash +# FusedAddRmsNorm — single-operator performance test (msprof op). +# +# Measures the student's operator only. NO correctness check here — use run.sh +# for that. +# +# Performance is measured on case 2 only (B=256, H=1024; one-based case ID). +# Usage (from anywhere — the script resolves the kit root itself): +# bash checker/profile.sh # Ascend C, case 2 +# LANG=triton bash checker/profile.sh # Triton, case 2 +# LANG=tilelang bash checker/profile.sh # TileLang, case 2 +set -e +ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)" +cd "$ROOT" + +if (( $# != 0 )); then + echo "[ERROR] profile.sh always profiles public case 2; no case argument is accepted." >&2 + exit 2 +fi +# Every `hpc submit` starts a fresh container, so always activate the complete +# CANN/Python/custom-OPP environment in this shell. +# shellcheck disable=SC1091 +source "$ROOT/env.sh" + +LANG="${LANG:-ascendc}" +# Resolve the backend. `LANG` doubles as the POSIX locale var on locale-enabled +# images (here it's exported as zh_CN.UTF-8), so a strict `== ascendc` check +# would wrongly skip the Ascend C build. See checker/run.sh for the rationale. +case "$LANG" in + triton|tilelang) BACKEND="$LANG"; NEED_BUILD=0 ;; + *) BACKEND="ascendc"; NEED_BUILD=1 ;; +esac +export PYTHONPATH="$ROOT:$ROOT/checker:${PYTHONPATH:-}" +PROFILE_CASE_NUM=2 +PROFILE_SHAPE="256x1024" +PYTHON_BIN="$(python3 -c 'import sys; print(sys.executable)')" + +# Ascend C needs the op + wheel installed in the current HPC container. A wheel +# left in the shared directory by an earlier job does not imply that its Python +# extension is installed in this fresh container. +if (( NEED_BUILD )); then + if ! python3 -c "import custom_ops_lib" >/dev/null 2>&1; then + bash checker/build.sh + fi + # build.sh runs in a subshell, so re-source the generated custom OPP env. + # shellcheck disable=SC1091 + source "$ROOT/env.sh" +fi + +echo "############################## STUDENT OP ($BACKEND) ##############################" +echo "=== [profile] student:$BACKEND case $PROFILE_CASE_NUM ($PROFILE_SHAPE) under msprof op ===" +PROF_DIR="prof_out" +rm -rf "$PROF_DIR" +mkdir -p "$PROF_DIR" +if ! LANG="$BACKEND" timeout 180 msprof op \ + --warm-up=10 \ + --output="$ROOT/$PROF_DIR" \ + "$PYTHON_BIN" checker/test_op.py --profile "$PROFILE_CASE_NUM"; then + echo "[ERROR] msprof op failed for student:$BACKEND" >&2 + exit 1 +fi +if ! STUDENT_US=$(python3 checker/get_time.py "$ROOT/$PROF_DIR"); then + echo "[ERROR] failed to read Task Duration for student:$BACKEND" >&2 + exit 1 +fi + +echo "" +echo "=========================== SUMMARY ===========================" +printf " student (%s): %s us\n" "$BACKEND" "$STUDENT_US" + +if [[ "$STUDENT_US" == "0.0000" ]]; then + echo " [ERROR] student op reported zero Task Duration" + exit 1 +fi +echo "Performance profiling complete." diff --git a/src/lab3p5/checker/run.sh b/src/lab3p5/checker/run.sh new file mode 100755 index 00000000..134c1a7a --- /dev/null +++ b/src/lab3p5/checker/run.sh @@ -0,0 +1,51 @@ +#!/bin/bash +# FusedAddRmsNorm — functional / correctness test (NO profiling). +# +# Runs `checker/test_op.py` (generates inputs from seed, fp32 golden, compares). +# Use `bash checker/profile.sh` for msprof performance numbers. +# +# Usage (from anywhere — the script resolves the kit root itself): +# bash checker/run.sh # Ascend C: all cases +# LANG=triton bash checker/run.sh # Triton: all cases (no build) +# LANG=tilelang bash checker/run.sh # TileLang: all cases (no build) +# bash checker/run.sh # single case +set -e +ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)" +cd "$ROOT" +# Every `hpc submit` starts a fresh container, so always activate the complete +# CANN/Python/custom-OPP environment in this shell. +# shellcheck disable=SC1091 +source "$ROOT/env.sh" + +LANG="${LANG:-ascendc}" +# Resolve the backend. `LANG` doubles as the POSIX locale var on locale-enabled +# images (here it's exported as zh_CN.UTF-8), so a strict `== ascendc` check +# would wrongly skip the Ascend C build. test_op.py only special-cases +# triton/tilelang and treats everything else as ascendc — mirror that: only +# those two values take the no-build path; any other value (incl. a locale) +# is the Ascend C default. This keeps `LANG=triton bash run.sh` (README) working. +case "$LANG" in + triton|tilelang) BACKEND="$LANG"; NEED_BUILD=0 ;; + *) BACKEND="ascendc"; NEED_BUILD=1 ;; +esac +export PYTHONPATH="$ROOT:$ROOT/checker:${PYTHONPATH:-}" +CASE_NUM="${1:-}" + +# Ascend C needs the op + wheel installed in the current HPC container. A wheel +# left in the shared directory by an earlier job is not enough. +if (( NEED_BUILD )); then + if ! python3 -c "import custom_ops_lib" >/dev/null 2>&1; then + bash checker/build.sh + fi + # build.sh runs in a subshell, so re-source the generated custom OPP env. + # shellcheck disable=SC1091 + source "$ROOT/env.sh" +fi + +echo "=== [run] correctness ($BACKEND) ===" +if [[ -n "$CASE_NUM" ]]; then + python3 -u checker/test_op.py "$CASE_NUM" +else + python3 -u checker/test_op.py +fi +echo "=== [run] correctness done ===" diff --git a/src/lab3p5/checker/test_op.py b/src/lab3p5/checker/test_op.py new file mode 100644 index 00000000..02d23df0 --- /dev/null +++ b/src/lab3p5/checker/test_op.py @@ -0,0 +1,180 @@ +#!/usr/bin/env python3 +# coding=utf-8 +"""FusedAddRmsNorm — correctness and single-launch profiling entry. + +Cases are generated on the fly from `case_specs.py` (random inputs by seed) — +NO .bin data ships in the kit. The judge swaps `case_specs.py` for a hidden one +with different shapes; this file itself never changes. + +Flow (mirrors the S9 samples Concat/Transpose): + 1. import custom_ops_lib (the student's built extension / Triton / TileLang) + 2. for each case: generate inputs from seed, run the op on NPU, compare + against the fp32-computed golden + +Profiling mode launches the selected operator exactly once. ``msprof op`` owns +the warm-up/replay loop and writes the final Task Duration to OpBasicInfo.csv. + +Usage (from kit root, after `source ./env.sh`): + python3 checker/test_op.py # run all cases + python3 checker/test_op.py # run a single case + python3 checker/test_op.py --profile # one launch, no verification +""" +import argparse +import os +import sys +import torch +import torch_npu + +# Pick the backend: Ascend C (default, via the pybind wheel), Triton +# (LANG=triton), or TileLang (LANG=tilelang). The latter two are pure Python and +# need no wheel build. +_lang = os.environ.get("LANG", "ascendc").lower() +if _lang in ("triton", "tilelang"): + import importlib + importlib.import_module(f"src.{_lang}") # registers `custom_ops_lib` in sys.modules + import custom_ops_lib +else: + import custom_ops_lib + +torch.npu.config.allow_internal_format = False +torch.npu.set_device(int(os.environ.get("ASCEND_DEVICE_ID", "0"))) + + +# --------------------------------------------------------------------------- +# Case specs (public; judge replaces this module with a hidden one) +# --------------------------------------------------------------------------- +import case_specs # noqa: E402 + + +DATA_RANGES = {"S": (-1.0, 1.0), "M": (1.0, 10.0), "L": (-1000.0, 1000.0)} + + +def gen_inputs(B, H, data_range, seed): + """Deterministic random inputs from seed (reproducible across runs/judge).""" + g = torch.Generator(device="cpu") + g.manual_seed(seed) + lo, hi = DATA_RANGES[data_range] + x = (torch.rand(B, H, generator=g) * (hi - lo) + lo).to(torch.float16) + r = (torch.rand(B, H, generator=g) * (hi - lo) + lo).to(torch.float16) + w = (torch.rand(H, generator=g) * 2.0).clamp(min=0.01).to(torch.float16) + return x, r, w + + +def golden(x, residual, weight, eps): + """fp32 compute then cast fp16 (matches the lab3.5_frame golden).""" + R = residual.float() + x.float() + ms = torch.mean(R * R, dim=-1, keepdim=True) + rms = torch.sqrt(ms + eps) + Y = (R / rms) * weight.float() + return Y.to(torch.float16), R.to(torch.float16) + + +# --------------------------------------------------------------------------- +# Verification (ported verbatim from the upstream checker verify_result.py) +# --------------------------------------------------------------------------- +def verify_result(real, golden_t): + """Element passes if abs_err<=tol OR rel_err<=tol, where rel_err uses + |golden| (clamped to eps) as the denominator — NOT max(|real|,|golden|). + The whole tensor passes only if error_ratio <= 0.0 (zero mismatches), + matching the upstream checker exactly. + """ + if golden_t.dtype == torch.float32: + tol = 1e-4 + else: + tol = 1e-3 + # Match the upstream checker: rel_err = abs_err / |golden| (clamped to eps). + out = real.detach().cpu().to(torch.float64).reshape(-1) + g = golden_t.detach().cpu().to(torch.float64).reshape(-1) + eps = 1e-12 + denom = torch.where(g.abs() < eps, torch.tensor(eps), g.abs()) + abs_err = (out - g).abs() + rel_err = abs_err / denom + pass_check = (abs_err <= tol) | (rel_err <= tol) + error_ratio = float((~pass_check).sum().item()) / g.numel() + # upstream: return error_ratio <= 0.0 + ok = error_ratio <= 0.0 + if not ok: + # show up to 100 failing indices, like the upstream checker + bad = torch.where(~pass_check)[0] + for i, idx in enumerate(bad[:100]): + gv = float(g[idx]); ov = float(out[idx]) + dv = abs(gv) if abs(gv) > eps else eps + print(f" idx={int(idx):06d} expected={gv:.9f} actual={ov:.9f} " + f"rdiff={abs(ov-gv)/dv:.6f}") + print(f" error_ratio={error_ratio:.6f} (tolerance: 0.0000)") + else: + print("test pass") + return ok + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- +def run_case(idx, spec): + B, H, eps, data_range, seed = spec + name = f"case_{idx}" + x, r, w = gen_inputs(B, H, data_range, seed) + gy, gres = golden(x, r, w, eps) # CPU golden + + x_d, r_d, w_d = x.npu(), r.npu(), w.npu() + y, resout = custom_ops_lib.fused_add_rmsnorm(x_d, r_d, w_d, eps) + if y is None: + print(f"[{name}] execution returned None (timeout?)") + return False + y = y.cpu() + resout = resout.cpu() + + ok_y = verify_result(y, gy) + ok_r = verify_result(resout, gres) + print(f"[{name}] {B}x{H} y={'pass' if ok_y else 'FAIL'} " + f"residual_out={'pass' if ok_r else 'FAIL'}") + return ok_y and ok_r + + +def profile_case(idx, spec): + """Launch only the selected op; msprof op performs warm-up and replay.""" + B, H, eps, data_range, seed = spec + x, r, w = gen_inputs(B, H, data_range, seed) + x_d, r_d, w_d = x.npu(), r.npu(), w.npu() + custom_ops_lib.fused_add_rmsnorm(x_d, r_d, w_d, eps) + torch.npu.synchronize() + print(f"[profile] case_{idx} {B}x{H} launched once") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("case_num", nargs="?", type=int) + parser.add_argument( + "--profile", + action="store_true", + help="launch one case once for msprof op; skip golden and verification", + ) + args = parser.parse_args() + + case_num = args.case_num + cases = case_specs.CASES + + if case_num is not None: + # 1-based case number, mirroring S9 test_op.py + i = case_num - 1 + if i < 0 or i >= len(cases): + print(f"[ERROR] case_{case_num} not found (have {len(cases)} cases)") + sys.exit(2) + if args.profile: + profile_case(i, cases[i]) + return + ok = run_case(i, cases[i]) + sys.exit(0 if ok else 1) + + if args.profile: + print("[ERROR] --profile requires a case number") + sys.exit(2) + + all_ok = True + for i, spec in enumerate(cases): + all_ok &= run_case(i, spec) + sys.exit(0 if all_ok else 1) + + +if __name__ == "__main__": + main() diff --git a/src/lab3p5/env.sh b/src/lab3p5/env.sh new file mode 100644 index 00000000..d4275de4 --- /dev/null +++ b/src/lab3p5/env.sh @@ -0,0 +1,81 @@ +#!/usr/bin/env bash +# lab3.5 student kit — environment activation. +# Source from the kit root: source ./env.sh +# +# Activates the system CANN toolkit and ensures python3 is on PATH. Relies only +# on system-level paths under /usr/local/Ascend and /usr/local/python* — no +# hard-coded home directories, no competition-specific conda env. +# +# NOTE: this file is *sourced* into an interactive shell, so do NOT add `set -e` +# here — it would leak into the caller's shell and terminate it on the first +# non-zero return. + +_KIT_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" + +# --- 1. System CANN toolkit ------------------------------------------------- +# The image may expose `msprof` without the full runtime library paths (notably +# libhccl.so). Therefore command availability is not a sufficient activation +# check. Source set_env.sh once per shell, preferring the `cann` symlink. +if [[ -z "${_LAB3P5_CANN_ENV_SOURCED:-}" ]]; then + for _set_env in \ + /usr/local/Ascend/cann/set_env.sh \ + /usr/local/Ascend/ascend-toolkit/latest/set_env.sh \ + /usr/local/Ascend/cann-8.5.0/set_env.sh; do + if [[ -f "$_set_env" ]]; then + # shellcheck disable=SC1090 + source "$_set_env" + _LAB3P5_CANN_ENV_SOURCED=1 + break + fi + done + unset _set_env +fi + +# --- 2. Logical NPU device --------------------------------------------------- +export ASCEND_DEVICE_ID="${ASCEND_DEVICE_ID:-0}" + +# --- 3. python3 -------------------------------------------------------------- +# The system interpreter lives outside the default PATH on this image; put it +# back if python3 isn't already resolvable. The kit depends on torch/torch_npu, +# which live in the cp311 site-packages under /usr/local/python3.11.14 — so prefer +# that interpreter before falling back to /usr/bin/python3 (the OS 3.10, which +# has none of the kit packages). /usr/local/bin is kept for cmake/ninja/pip. +if ! command -v python3 >/dev/null 2>&1 || ! python3 -c "import torch, torch_npu" >/dev/null 2>&1; then + for _py in /usr/local/python3.11.14/bin /usr/local/bin /usr/bin; do + if [[ -x "$_py/python3" ]]; then + export PATH="$_py:$PATH" + break + fi + done + unset _py +fi + +# --- 4. Custom operator package (built by build.sh) ------------------------- +# set_env.bash is generated at build time and already exports +# ASCEND_CUSTOM_OPP_PATH plus prepends op_api/lib to LD_LIBRARY_PATH — so we +# only source it. Guard with a flag so re-sourcing doesn't pile up duplicates. +# CUSTOM_OPP_HOME matches build.sh's --install-path (overridable). +export CUSTOM_OPP_HOME="${CUSTOM_OPP_HOME:-$HOME/custom_opp}" +if [[ -z "${CUSTOM_OPP_SOURCED:-}" ]] && [[ -f "$CUSTOM_OPP_HOME/vendors/customize/bin/set_env.bash" ]]; then + # shellcheck disable=SC1091 + source "$CUSTOM_OPP_HOME/vendors/customize/bin/set_env.bash" + export CUSTOM_OPP_SOURCED=1 +fi + +# --- 5. Make the kit importable -------------------------------------------- +# test_op.py imports `case_specs`, `custom_ops_lib`, and `src.` by name. +# `src` lives under the kit root; `case_specs`/`test_op`/`get_time` live under +# the sibling `checker/` dir. Put both on PYTHONPATH so a bare +# `python3 checker/test_op.py` resolves every import. +export PYTHONPATH="$_KIT_ROOT:$_KIT_ROOT/checker${PYTHONPATH:+:$PYTHONPATH}" + +# --- 6. Sanity checks (loud, not silent) ----------------------------------- +if [[ -z "${ASCEND_TOOLKIT_HOME:-}" ]]; then + echo "[env.sh] ERROR: ASCEND_TOOLKIT_HOME is unset — no CANN set_env.sh found under /usr/local/Ascend" >&2 +fi +if ! command -v python3 >/dev/null 2>&1; then + echo "[env.sh] ERROR: python3 not found on PATH" >&2 +fi + +echo "[env.sh] CANN=${ASCEND_TOOLKIT_HOME:-?} device=${ASCEND_DEVICE_ID} python=$(command -v python3 || echo '?')" +unset _KIT_ROOT diff --git a/src/lab3p5/src/__init__.py b/src/lab3p5/src/__init__.py new file mode 100644 index 00000000..208a5474 --- /dev/null +++ b/src/lab3p5/src/__init__.py @@ -0,0 +1 @@ +"""Marker so `src` is a package and `src.` is importable.""" diff --git a/src/lab3p5/src/ascendc/CMakeLists.txt b/src/lab3p5/src/ascendc/CMakeLists.txt new file mode 100644 index 00000000..4e82197c --- /dev/null +++ b/src/lab3p5/src/ascendc/CMakeLists.txt @@ -0,0 +1,31 @@ +# ---------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ---------------------------------------------------------------------------------------------------------- + + +cmake_minimum_required(VERSION 3.16.0) +project(opp) +find_package(ASC REQUIRED) +set(package_name ${vendor_name}) + +npu_op_package(${package_name} + TYPE RUN + CONFIG + INSTALL_PATH ${CMAKE_BINARY_DIR}/ +) + +if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/framework) + add_subdirectory(framework) +endif() +if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/op_host) + add_subdirectory(op_host) +endif() +if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/op_kernel) + add_subdirectory(op_kernel) +endif() \ No newline at end of file diff --git a/src/lab3p5/src/ascendc/CMakePresets.json b/src/lab3p5/src/ascendc/CMakePresets.json new file mode 100644 index 00000000..21212bb5 --- /dev/null +++ b/src/lab3p5/src/ascendc/CMakePresets.json @@ -0,0 +1,67 @@ +{ + "version": 1, + "cmakeMinimumRequired": { + "major": 3, + "minor": 16, + "patch": 0 + }, + "configurePresets": [ + { + "name": "default", + "displayName": "Default Config", + "description": "Default build using Unix Makefiles generator", + "generator": "Unix Makefiles", + "binaryDir": "${sourceDir}/build_out", + "cacheVariables": { + "CMAKE_BUILD_TYPE": { + "type": "STRING", + "value": "Release" + }, + "ENABLE_SOURCE_PACKAGE": { + "type": "BOOL", + "value": "True" + }, + "ENABLE_BINARY_PACKAGE": { + "type": "BOOL", + "value": "True" + }, + "ASCEND_COMPUTE_UNIT": { + "type": "STRING", + "value": "ascend910b" + }, + "ENABLE_TEST": { + "type": "BOOL", + "value": "True" + }, + "vendor_name": { + "type": "STRING", + "value": "customize" + }, + "ASCEND_CANN_PACKAGE_PATH": { + "type": "PATH", + "value": "/usr/local/Ascend/cann-8.5.0" + }, + "ASCEND_PYTHON_EXECUTABLE": { + "type": "STRING", + "value": "python3" + }, + "CMAKE_INSTALL_PREFIX": { + "type": "PATH", + "value": "${sourceDir}/build_out" + }, + "ENABLE_CROSS_COMPILE": { + "type": "BOOL", + "value": "False" + }, + "CMAKE_CROSS_PLATFORM_COMPILER": { + "type": "PATH", + "value": "/usr/bin/aarch64-linux-gnu-g++" + }, + "ASCEND_PACK_SHARED_LIBRARY": { + "type": "BOOL", + "value": "False" + } + } + } + ] +} \ No newline at end of file diff --git a/src/lab3p5/src/ascendc/build_op.sh b/src/lab3p5/src/ascendc/build_op.sh new file mode 100755 index 00000000..9759204f --- /dev/null +++ b/src/lab3p5/src/ascendc/build_op.sh @@ -0,0 +1,39 @@ +#!/bin/bash +if [ -z "$BASE_LIBS_PATH" ]; then + if [ -z "$ASCEND_HOME_PATH" ]; then + if [ -z "$ASCEND_AICPU_PATH" ]; then + echo "please set env." + exit 1 + else + export ASCEND_HOME_PATH=$ASCEND_AICPU_PATH + fi + else + export ASCEND_HOME_PATH=$ASCEND_HOME_PATH + fi +else + export ASCEND_HOME_PATH=$BASE_LIBS_PATH +fi +echo "using ASCEND_HOME_PATH: $ASCEND_HOME_PATH" +script_path=$(realpath $(dirname $0)) + +BUILD_DIR="build_out" +HOST_NATIVE_DIR="host_native_tiling" +mkdir -p build_out +rm -rf build_out/* + +ENABLE_CROSS="-DENABLE_CROSS_COMPILE=True" +ENABLE_BINARY="-DENABLE_BINARY_PACKAGE=True" +ENABLE_LIBRARY="-DASCEND_PACK_SHARED_LIBRARY=True" +cmake_version=$(cmake --version | grep "cmake version" | awk '{print $3}') + +target=package +if [ "$1"x != ""x ]; then target=$1; fi + +cmake -S . -B "$BUILD_DIR" --preset=default +cmake --build "$BUILD_DIR" --target binary -j$(nproc) +cmake --build "$BUILD_DIR" --target $target -j$(nproc) + +# Install path mirrors build.sh / env.sh (CUSTOM_OPP_HOME, plural). Falling back +# to $HOME/custom_opp keeps a bare `bash build_op.sh` self-consistent. +INSTALL_PATH="${CUSTOM_OPP_HOME:-$HOME/custom_opp}" +"$BUILD_DIR/custom_opp_ubuntu_aarch64.run" --quiet --install-path="$INSTALL_PATH" \ No newline at end of file diff --git a/src/lab3p5/src/ascendc/common/pytorch_npu_helper.hpp b/src/lab3p5/src/ascendc/common/pytorch_npu_helper.hpp new file mode 100644 index 00000000..6414dfa5 --- /dev/null +++ b/src/lab3p5/src/ascendc/common/pytorch_npu_helper.hpp @@ -0,0 +1,583 @@ +/****************************************************************************** + * Copyright (c) 2022 Huawei Technologies Co., Ltd + * All rights reserved. + * + * Licensed under the BSD 3-Clause License (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://opensource.org/licenses/BSD-3-Clause + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + ******************************************************************************/ + +#ifndef PYTORCH_NPU_HELPER_HPP_ +#define PYTORCH_NPU_HELPER_HPP_ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + + +#include "torch_npu/csrc/aten/NPUNativeFunctions.h" +#include "torch_npu/csrc/core/npu/NPUStream.h" +#include "torch_npu/csrc/framework/OpCommand.h" +#include "torch_npu/csrc/framework/interface/EnvVariables.h" +#include "torch_npu/csrc/framework/utils/CalcuOpUtil.h" +#include "torch_npu/csrc/framework/utils/OpPreparation.h" + +#define NPU_NAME_SPACE at_npu::native + +#define __FILENAME__ (strrchr("/" __FILE__, '/') + 1) + +typedef struct aclOpExecutor aclOpExecutor; +typedef struct aclTensor aclTensor; +typedef struct aclScalar aclScalar; +typedef struct aclIntArray aclIntArray; +typedef struct aclFloatArray aclFloatArray; +typedef struct aclBoolArray aclBoolArray; +typedef struct aclTensorList aclTensorList; + +typedef aclTensor *(*_aclCreateTensor)( + const int64_t *view_dims, uint64_t view_dims_num, aclDataType data_type, + const int64_t *stride, int64_t offset, aclFormat format, + const int64_t *storage_dims, uint64_t storage_dims_num, void *tensor_data); +typedef aclScalar *(*_aclCreateScalar)(void *value, aclDataType data_type); +typedef aclIntArray *(*_aclCreateIntArray)(const int64_t *value, uint64_t size); +typedef aclFloatArray *(*_aclCreateFloatArray)(const float *value, + uint64_t size); +typedef aclBoolArray *(*_aclCreateBoolArray)(const bool *value, uint64_t size); +typedef aclTensorList *(*_aclCreateTensorList)(const aclTensor *const *value, + uint64_t size); + +typedef int (*_aclDestroyTensor)(const aclTensor *tensor); +typedef int (*_aclDestroyScalar)(const aclScalar *scalar); +typedef int (*_aclDestroyIntArray)(const aclIntArray *array); +typedef int (*_aclDestroyFloatArray)(const aclFloatArray *array); +typedef int (*_aclDestroyBoolArray)(const aclBoolArray *array); +typedef int (*_aclDestroyTensorList)(const aclTensorList *array); + +constexpr int kHashBufSize = 8192; +constexpr int kHashBufMaxSize = kHashBufSize + 1024; +extern thread_local char g_hashBuf[kHashBufSize]; +extern thread_local int g_hashOffset; + + +#define AT_ALL_SCALAR_TYPE_AND_ACL_DATATYPE_PAIR(_) \ + _(at::ScalarType::Byte, ACL_UINT8) \ + _(at::ScalarType::Char, ACL_INT8) \ + _(at::ScalarType::Short, ACL_INT16) \ + _(at::ScalarType::Int, ACL_INT32) \ + _(at::ScalarType::Long, ACL_INT64) \ + _(at::ScalarType::Half, ACL_FLOAT16) \ + _(at::ScalarType::Float, ACL_FLOAT) \ + _(at::ScalarType::Double, ACL_DOUBLE) \ + _(at::ScalarType::ComplexHalf, ACL_DT_UNDEFINED) \ + _(at::ScalarType::ComplexFloat, ACL_COMPLEX64) \ + _(at::ScalarType::ComplexDouble, ACL_COMPLEX128) \ + _(at::ScalarType::Bool, ACL_BOOL) \ + _(at::ScalarType::QInt8, ACL_DT_UNDEFINED) \ + _(at::ScalarType::QUInt8, ACL_DT_UNDEFINED) \ + _(at::ScalarType::QInt32, ACL_DT_UNDEFINED) \ + _(at::ScalarType::BFloat16, ACL_BF16) \ + _(at::ScalarType::QUInt4x2, ACL_DT_UNDEFINED) \ + _(at::ScalarType::QUInt2x4, ACL_DT_UNDEFINED) \ + _(at::ScalarType::Undefined, ACL_DT_UNDEFINED) \ + _(at::ScalarType::NumOptions, ACL_DT_UNDEFINED) + +constexpr aclDataType kATenScalarTypeToAclDataTypeTable + [static_cast(at::ScalarType::NumOptions) + 1] = { +#define DEFINE_ENUM(_1, n) n, + AT_ALL_SCALAR_TYPE_AND_ACL_DATATYPE_PAIR(DEFINE_ENUM) +#undef DEFINE_ENUM +}; + +#define GET_OP_API_FUNC(apiName) \ + reinterpret_cast<_##apiName>(GetOpApiFuncAddr(#apiName)) + +#define MEMCPY_TO_BUF(data_expression, size_expression) \ + if (g_hashOffset + (size_expression) > kHashBufSize) { \ + g_hashOffset = kHashBufMaxSize; \ + return; \ + } \ + memcpy(g_hashBuf + g_hashOffset, data_expression, size_expression); \ + g_hashOffset += size_expression; + +inline const char *GetOpApiLibName(void) { return "libopapi.so"; } + +inline const char *GetCustOpApiLibName(void) { return "libcust_opapi.so"; } + +inline void *GetOpApiFuncAddrInLib(void *handler, const char *libName, + const char *apiName) { + auto funcAddr = dlsym(handler, apiName); + if (funcAddr == nullptr) { + ASCEND_LOGW("dlsym %s from %s failed, error:%s.", apiName, libName, + dlerror()); + } + return funcAddr; +} + +inline void *GetOpApiLibHandler(const char *libName) { + auto handler = dlopen(libName, RTLD_LAZY); + if (handler == nullptr) { + ASCEND_LOGW("dlopen %s failed, error:%s.", libName, dlerror()); + } + return handler; +} + +inline void *GetOpApiFuncAddr(const char *apiName) { + static auto custOpApiHandler = GetOpApiLibHandler(GetCustOpApiLibName()); + if (custOpApiHandler != nullptr) { + auto funcAddr = + GetOpApiFuncAddrInLib(custOpApiHandler, GetCustOpApiLibName(), apiName); + if (funcAddr != nullptr) { + return funcAddr; + } + } + + static auto opApiHandler = GetOpApiLibHandler(GetOpApiLibName()); + if (opApiHandler == nullptr) { + return nullptr; + } + return GetOpApiFuncAddrInLib(opApiHandler, GetOpApiLibName(), apiName); +} + +inline c10::Scalar ConvertTensorToScalar(const at::Tensor &tensor) { + c10::Scalar expScalar; + const at::Tensor *aclInput = &tensor; + if (aclInput->scalar_type() == at::ScalarType::Double) { + double value = *(double *)aclInput->data_ptr(); + c10::Scalar scalar(value); + expScalar = scalar; + } else if (aclInput->scalar_type() == at::ScalarType::Long) { + int64_t value = *(int64_t *)aclInput->data_ptr(); + c10::Scalar scalar(value); + expScalar = scalar; + } else if (aclInput->scalar_type() == at::ScalarType::Float) { + float value = *(float *)aclInput->data_ptr(); + c10::Scalar scalar(value); + expScalar = scalar; + } else if (aclInput->scalar_type() == at::ScalarType::Int) { + int value = *(int *)aclInput->data_ptr(); + c10::Scalar scalar(value); + expScalar = scalar; + } else if (aclInput->scalar_type() == at::ScalarType::Half) { + c10::Half value = *(c10::Half *)aclInput->data_ptr(); + c10::Scalar scalar(value); + expScalar = scalar; + } else if (aclInput->scalar_type() == at::ScalarType::Bool) { + int8_t value = *(int8_t *)aclInput->data_ptr(); + c10::Scalar scalar(value); + expScalar = scalar; + } else if (aclInput->scalar_type() == at::ScalarType::ComplexDouble) { + c10::complex value = *(c10::complex *)aclInput->data_ptr(); + c10::Scalar scalar(value); + expScalar = scalar; + } else if (aclInput->scalar_type() == at::ScalarType::ComplexFloat) { + c10::complex value = *(c10::complex *)aclInput->data_ptr(); + c10::Scalar scalar(value); + expScalar = scalar; + } else if (aclInput->scalar_type() == at::ScalarType::BFloat16) { + c10::BFloat16 value = *(c10::BFloat16 *)aclInput->data_ptr(); + c10::Scalar scalar(value); + expScalar = scalar; + } + return expScalar; +} + +inline at::Tensor CopyTensorHostToDevice(const at::Tensor &cpu_tensor) { + at::Tensor cpuPinMemTensor = cpu_tensor.pin_memory(); + int deviceIndex = 0; + return cpuPinMemTensor.to(c10::Device(torch_npu::utils::get_npu_device_type(), deviceIndex), + cpuPinMemTensor.scalar_type(), true, true); +} + +inline at::Tensor CopyScalarToDevice(const c10::Scalar &cpu_scalar, + at::ScalarType scalar_data_type) { + return CopyTensorHostToDevice( + scalar_to_tensor(cpu_scalar).to(scalar_data_type)); +} + +inline aclTensor *ConvertType(const at::Tensor &at_tensor) { + static const auto aclCreateTensor = GET_OP_API_FUNC(aclCreateTensor); + if (aclCreateTensor == nullptr) { + return nullptr; + } + + if (!at_tensor.defined()) { + return nullptr; + } + at::ScalarType scalar_data_type = at_tensor.scalar_type(); + aclDataType acl_data_type = + kATenScalarTypeToAclDataTypeTable[static_cast(scalar_data_type)]; + TORCH_CHECK( + acl_data_type != ACL_DT_UNDEFINED, + std::string(c10::toString(scalar_data_type)) + " has not been supported") + c10::SmallVector storageDims; + // if acl_data_type is ACL_STRING, storageDims is empty. + auto itemsize = at_tensor.itemsize(); + if (itemsize == 0) { + AT_ERROR("When ConvertType, tensor item size of cannot be zero."); + return nullptr; + } + if (acl_data_type != ACL_STRING) { + storageDims.push_back(at_tensor.storage().nbytes() / itemsize); + } + + const auto dimNum = at_tensor.sizes().size(); + aclFormat format = ACL_FORMAT_ND; + switch (dimNum) { + case 3: + format = ACL_FORMAT_NCL; + break; + case 4: + format = ACL_FORMAT_NCHW; + break; + case 5: + format = ACL_FORMAT_NCDHW; + break; + default: + format = ACL_FORMAT_ND; + } + + if (at_tensor.unsafeGetTensorImpl()->is_wrapped_number()) { + c10::Scalar expScalar = ConvertTensorToScalar(at_tensor); + at::Tensor aclInput = CopyScalarToDevice(expScalar, scalar_data_type); + return aclCreateTensor(aclInput.sizes().data(), aclInput.sizes().size(), + acl_data_type, aclInput.strides().data(), + aclInput.storage_offset(), format, + storageDims.data(), storageDims.size(), + const_cast(aclInput.storage().data())); + } + + auto acl_tensor = aclCreateTensor( + at_tensor.sizes().data(), at_tensor.sizes().size(), acl_data_type, + at_tensor.strides().data(), at_tensor.storage_offset(), format, + storageDims.data(), storageDims.size(), + const_cast(at_tensor.storage().data())); + return acl_tensor; +} + +inline aclScalar *ConvertType(const at::Scalar &at_scalar) { + static const auto aclCreateScalar = GET_OP_API_FUNC(aclCreateScalar); + if (aclCreateScalar == nullptr) { + return nullptr; + } + + at::ScalarType scalar_data_type = at_scalar.type(); + aclDataType acl_data_type = + kATenScalarTypeToAclDataTypeTable[static_cast(scalar_data_type)]; + TORCH_CHECK( + acl_data_type != ACL_DT_UNDEFINED, + std::string(c10::toString(scalar_data_type)) + " has not been supported") + aclScalar *acl_scalar = nullptr; + switch (scalar_data_type) { + case at::ScalarType::Double: { + double value = at_scalar.toDouble(); + acl_scalar = aclCreateScalar(&value, acl_data_type); + break; + } + case at::ScalarType::Long: { + int64_t value = at_scalar.toLong(); + acl_scalar = aclCreateScalar(&value, acl_data_type); + break; + } + case at::ScalarType::Bool: { + bool value = at_scalar.toBool(); + acl_scalar = aclCreateScalar(&value, acl_data_type); + break; + } + case at::ScalarType::ComplexDouble: { + auto value = at_scalar.toComplexDouble(); + acl_scalar = aclCreateScalar(&value, acl_data_type); + break; + } + default: + acl_scalar = nullptr; + break; + } + return acl_scalar; +} + +inline aclIntArray *ConvertType(const at::IntArrayRef &at_array) { + static const auto aclCreateIntArray = GET_OP_API_FUNC(aclCreateIntArray); + if (aclCreateIntArray == nullptr) { + return nullptr; + } + auto array = aclCreateIntArray(at_array.data(), at_array.size()); + return array; +} + +template +inline aclBoolArray *ConvertType(const std::array &value) { + static const auto aclCreateBoolArray = GET_OP_API_FUNC(aclCreateBoolArray); + if (aclCreateBoolArray == nullptr) { + return nullptr; + } + + auto array = aclCreateBoolArray(value.data(), value.size()); + return array; +} + +inline aclBoolArray *ConvertType(const at::ArrayRef &value) { + static const auto aclCreateBoolArray = GET_OP_API_FUNC(aclCreateBoolArray); + if (aclCreateBoolArray == nullptr) { + return nullptr; + } + + auto array = aclCreateBoolArray(value.data(), value.size()); + return array; +} + +inline aclTensorList *ConvertType(const at::TensorList &at_tensor_list) { + static const auto aclCreateTensorList = GET_OP_API_FUNC(aclCreateTensorList); + if (aclCreateTensorList == nullptr) { + return nullptr; + } + + std::vector tensor_list(at_tensor_list.size()); + for (size_t i = 0; i < at_tensor_list.size(); i++) { + tensor_list[i] = ConvertType(at_tensor_list[i]); + } + auto acl_tensor_list = + aclCreateTensorList(tensor_list.data(), tensor_list.size()); + return acl_tensor_list; +} + +inline aclTensor *ConvertType(const c10::optional &opt_tensor) { + if (opt_tensor.has_value() && opt_tensor.value().defined()) { + return ConvertType(opt_tensor.value()); + } + return nullptr; +} + +inline aclIntArray *ConvertType( + const c10::optional &opt_array) { + if (opt_array.has_value()) { + return ConvertType(opt_array.value()); + } + return nullptr; +} + +inline aclScalar *ConvertType(const c10::optional &opt_scalar) { + if (opt_scalar.has_value()) { + return ConvertType(opt_scalar.value()); + } + return nullptr; +} + +inline aclDataType ConvertType(const at::ScalarType scalarType) { + return kATenScalarTypeToAclDataTypeTable[static_cast(scalarType)]; +} + +template +T ConvertType(T value) { + return value; +} + +template +auto ConvertToOpApiFunc(const Tuple ¶ms, void *opApiAddr, + std::index_sequence) { + typedef int (*OpApiFunc)( + typename std::decay(params))>::type...); + auto func = reinterpret_cast(opApiAddr); + return func; +} + +template +auto ConvertToOpApiFunc(const Tuple ¶ms, void *opApiAddr) { + static constexpr auto size = std::tuple_size::value; + return ConvertToOpApiFunc(params, opApiAddr, + std::make_index_sequence{}); +} + +inline void Release(aclTensor *p) { + static const auto aclDestroyTensor = GET_OP_API_FUNC(aclDestroyTensor); + if (aclDestroyTensor == nullptr) { + return; + } + aclDestroyTensor(p); +} + +inline void Release(aclScalar *p) { + static const auto aclDestroyScalar = GET_OP_API_FUNC(aclDestroyScalar); + if (aclDestroyScalar == nullptr) { + return; + } + aclDestroyScalar(p); +} + +inline void Release(aclIntArray *p) { + static const auto aclDestroyIntArray = GET_OP_API_FUNC(aclDestroyIntArray); + if (aclDestroyIntArray == nullptr) { + return; + } + + aclDestroyIntArray(p); +} + +inline void Release(aclBoolArray *p) { + static const auto aclDestroyBoolArray = GET_OP_API_FUNC(aclDestroyBoolArray); + if (aclDestroyBoolArray == nullptr) { + return; + } + + aclDestroyBoolArray(p); +} + +inline void Release(aclTensorList *p) { + static const auto aclDestroyTensorList = + GET_OP_API_FUNC(aclDestroyTensorList); + if (aclDestroyTensorList == nullptr) { + return; + } + + aclDestroyTensorList(p); +} + +template +void Release(T value) { + (void)value; +} + +template +void CallRelease(Tuple t, std::index_sequence) { + (void)std::initializer_list{(Release(std::get(t)), 0)...}; +} + +template +void ReleaseConvertTypes(Tuple &t) { + static constexpr auto size = std::tuple_size::value; + CallRelease(t, std::make_index_sequence{}); +} + +template +constexpr auto ConvertTypes(Ts &... args) { + return std::make_tuple(ConvertType(args)...); +} + +template +auto call(Function f, Tuple t, std::index_sequence) { + return f(std::get(t)...); +} + +template +auto call(Function f, Tuple t) { + static constexpr auto size = std::tuple_size::value; + return call(f, t, std::make_index_sequence{}); +} + +template +void AddParamToBuf(const std::array &value) { + MEMCPY_TO_BUF(value.data(), value.size() * sizeof(bool)); +} + +template +void AddParamToBuf(const T &value) { + MEMCPY_TO_BUF(&value, sizeof(T)); +} + +void AddParamToBuf(const at::Tensor &); +void AddParamToBuf(const at::Scalar &); +void AddParamToBuf(const at::IntArrayRef &); +void AddParamToBuf(const at::ArrayRef &); +void AddParamToBuf(const at::TensorList &); +void AddParamToBuf(const c10::optional &); +void AddParamToBuf(const c10::optional &); +void AddParamToBuf(const c10::optional &); +void AddParamToBuf(const at::ScalarType); +void AddParamToBuf(const string &); +void AddParamToBuf(); + +template +void AddParamToBuf(const T &arg, Args &... args) { + AddParamToBuf(arg); + AddParamToBuf(args...); +} + +uint64_t CalcHashId(); +typedef int (*InitHugeMemThreadLocal)(void *, bool); +typedef void (*UnInitHugeMemThreadLocal)(void *, bool); +typedef void (*ReleaseHugeMem)(void *, bool); + +#define EXEC_NPU_CMD(aclnn_api, ...) \ + do { \ + static const auto getWorkspaceSizeFuncAddr = \ + GetOpApiFuncAddr(#aclnn_api "GetWorkspaceSize"); \ + static const auto opApiFuncAddr = GetOpApiFuncAddr(#aclnn_api); \ + static const auto initMemAddr = \ + GetOpApiFuncAddr("InitHugeMemThreadLocal"); \ + static const auto unInitMemAddr = \ + GetOpApiFuncAddr("UnInitHugeMemThreadLocal"); \ + static const auto releaseMemAddr = GetOpApiFuncAddr("ReleaseHugeMem"); \ + TORCH_CHECK( \ + getWorkspaceSizeFuncAddr != nullptr && opApiFuncAddr != nullptr, \ + #aclnn_api, " or ", #aclnn_api "GetWorkspaceSize", " not in ", \ + GetOpApiLibName(), ", or ", GetOpApiLibName(), "not found."); \ + auto acl_stream = c10_npu::getCurrentNPUStream().stream(false); \ + uint64_t workspace_size = 0; \ + uint64_t *workspace_size_addr = &workspace_size; \ + aclOpExecutor *executor = nullptr; \ + aclOpExecutor **executor_addr = &executor; \ + InitHugeMemThreadLocal initMemFunc = \ + reinterpret_cast(initMemAddr); \ + UnInitHugeMemThreadLocal unInitMemFunc = \ + reinterpret_cast(unInitMemAddr); \ + if (initMemFunc) { \ + initMemFunc(nullptr, false); \ + } \ + auto converted_params = \ + ConvertTypes(__VA_ARGS__, workspace_size_addr, executor_addr); \ + static auto getWorkspaceSizeFunc = \ + ConvertToOpApiFunc(converted_params, getWorkspaceSizeFuncAddr); \ + auto workspace_status = call(getWorkspaceSizeFunc, converted_params); \ + TORCH_CHECK(workspace_status == 0, \ + "call " #aclnn_api " failed, detail:", aclGetRecentErrMsg()); \ + void *workspace_addr = nullptr; \ + if (workspace_size != 0) { \ + at::TensorOptions options = \ + at::TensorOptions(torch_npu::utils::get_npu_device_type()); \ + auto workspace_tensor = \ + at::empty({workspace_size}, options.dtype(kByte)); \ + workspace_addr = const_cast(workspace_tensor.storage().data()); \ + } \ + auto acl_call = [converted_params, workspace_addr, workspace_size, \ + acl_stream, executor]() -> int { \ + typedef int (*OpApiFunc)(void *, uint64_t, aclOpExecutor *, \ + const aclrtStream); \ + OpApiFunc opApiFunc = reinterpret_cast(opApiFuncAddr); \ + auto api_ret = \ + opApiFunc(workspace_addr, workspace_size, executor, acl_stream); \ + TORCH_CHECK(api_ret == 0, "call " #aclnn_api " failed, detail:", \ + aclGetRecentErrMsg()); \ + ReleaseConvertTypes(converted_params); \ + ReleaseHugeMem releaseMemFunc = \ + reinterpret_cast(releaseMemAddr); \ + if (releaseMemFunc) { \ + releaseMemFunc(nullptr, false); \ + } \ + return api_ret; \ + }; \ + at_npu::native::OpCommand cmd; \ + cmd.Name(#aclnn_api); \ + cmd.SetCustomHandler(acl_call); \ + cmd.Run(); \ + if (unInitMemFunc) { \ + unInitMemFunc(nullptr, false); \ + } \ + } while (false) + +#endif // PYTORCH_NPU_HELPER_HPP_ \ No newline at end of file diff --git a/src/lab3p5/src/ascendc/extension/custom_op.cpp b/src/lab3p5/src/ascendc/extension/custom_op.cpp new file mode 100644 index 00000000..6d803740 --- /dev/null +++ b/src/lab3p5/src/ascendc/extension/custom_op.cpp @@ -0,0 +1,62 @@ +/* + * FusedAddRmsNorm — Ascend C (aclnn) entry for the student framework. + * + * This file is the ONLY C++ the student must touch for the Ascend C backend. + * It wires the op's aclnn API into a torch custom op so test_op.py can call it. + * + * The actual kernel lives in op/op_host + op/op_kernel (the FusedAddRmsNorm + * aclnn op you build with build.sh). Here we just launch it via EXEC_NPU_CMD, + * exactly like the S9 samples (Concat/Transpose). + */ +#include +#include "../common/pytorch_npu_helper.hpp" + +using namespace at; + +/** + * FusedAddRmsNorm host launch. + * residual_out = residual + x + * y = residual_out / sqrt(mean(residual_out^2, dim=-1) + eps) * weight + * + * Args: + * x : (B, H) float16 + * residual : (B, H) float16 + * weight : (H,) float16 + * eps : float + * Returns: + * tuple(y, residual_out) — both (B, H) float16. + * + * NOTE: this default impl just calls the aclnn API you registered. If you + * prefer a hand-written Ascend C kernel (op/op_kernel/*.cpp), that path is + * wired too — just make sure aclnnFusedAddRmsNorm resolves to your kernel + * after build.sh + install. + */ +std::tuple fused_add_rmsnorm_impl_npu( + const Tensor &x, const Tensor &residual, const Tensor &weight, + double eps) { + Tensor y = at::empty_like(x); + Tensor residual_out = at::empty_like(x); + + // aclnnFusedAddRmsNorm(x, residual, weight, eps, enable_pdl, y, residual_out) + // Profiling warm-up is handled by `msprof op --warm-up`, so a normal op + // invocation launches exactly one FusedAddRmsNorm kernel. + bool enable_pdl = false; + EXEC_NPU_CMD(aclnnFusedAddRmsNorm, + x, residual, weight, eps, enable_pdl, + y, residual_out); + return std::make_tuple(y, residual_out); +} + +TORCH_LIBRARY(fusedaddrmsnorm, m) { + m.def("fused_add_rmsnorm(Tensor x, Tensor residual, Tensor weight, " + "float eps) -> (Tensor y, Tensor residual_out)"); +} + +TORCH_LIBRARY_IMPL(fusedaddrmsnorm, PrivateUse1, m) { + m.impl("fused_add_rmsnorm", &fused_add_rmsnorm_impl_npu); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("fused_add_rmsnorm", &fused_add_rmsnorm_impl_npu, + "FusedAddRmsNorm (residual+x then RMSNorm)"); +} diff --git a/src/lab3p5/src/ascendc/op_host/CMakeLists.txt b/src/lab3p5/src/ascendc/op_host/CMakeLists.txt new file mode 100644 index 00000000..d2be0a50 --- /dev/null +++ b/src/lab3p5/src/ascendc/op_host/CMakeLists.txt @@ -0,0 +1,57 @@ +# ---------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ---------------------------------------------------------------------------------------------------------- + + +aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR} ops_srcs) +npu_op_code_gen( + SRC ${ops_srcs} + PACKAGE ${package_name} + OUT_DIR ${ASCEND_AUTOGEN_PATH} +) + +file(GLOB autogen_aclnn_src ${ASCEND_AUTOGEN_PATH}/aclnn_*.cpp) +set_source_files_properties(${autogen_aclnn_src} PROPERTIES GENERATED TRUE) +npu_op_library(cust_opapi ACLNN + ${autogen_aclnn_src} +) + +target_compile_options(cust_opapi PRIVATE + -fvisibility=hidden +) + +file(GLOB group_proto_src ${ASCEND_AUTOGEN_PATH}/group_op_proto/*.cc) +file(GLOB proto_src ${ASCEND_AUTOGEN_PATH}/op_proto.cc) +set_source_files_properties(${group_proto_src} PROPERTIES GENERATED TRUE) +set_source_files_properties(${proto_src} PROPERTIES GENERATED TRUE) +npu_op_library(cust_op_proto GRAPH + ${ops_srcs} + ${group_proto_src} + ${proto_src} +) +target_compile_options(cust_op_proto PRIVATE + -fvisibility=hidden +) + +file(GLOB fallback_src ${ASCEND_AUTOGEN_PATH}/fallback_*.cpp) +set_source_files_properties(${fallback_src} PROPERTIES GENERATED TRUE) +npu_op_library(cust_optiling TILING + ${ops_srcs} + ${fallback_src} +) +target_compile_options(cust_op_proto PRIVATE + -fvisibility=hidden +) + +npu_op_package_add(${package_name} + LIBRARY + cust_optiling + cust_opapi + cust_op_proto +) diff --git a/src/lab3p5/src/ascendc/op_host/fused_add_rms_norm.cpp b/src/lab3p5/src/ascendc/op_host/fused_add_rms_norm.cpp new file mode 100644 index 00000000..7182ed91 --- /dev/null +++ b/src/lab3p5/src/ascendc/op_host/fused_add_rms_norm.cpp @@ -0,0 +1,152 @@ +/** + * @file fused_add_rms_norm.cpp + * @brief Host tiling + registration for the FusedAddRmsNorm operator. + * + * Op: FusedAddRmsNorm(x, residual, weight, eps, enable_pdl) -> (y, residual_out) + * residual_out = x + residual + * y = residual_out / sqrt(mean(residual_out^2, dim=-1) + eps) * weight + * + * Inputs are FP16 ND tensors of shape (B, H); weight is (H,). Outputs mirror the + * (B, H) shape. The kernel is row-parallel and reduces the H axis in FP32. + * + * Tiling carries B, H, the 32B-aligned H (for UB sizing), the align unit, and + * eps. Multi-core split is over the B (row) axis: each AIV core owns a disjoint + * contiguous range of rows — no cross-core sync or workspace needed. + */ +#include "../op_kernel/fused_add_rms_norm_tiling.h" +#include "register/op_def_registry.h" +#include "tiling/platform/platform_ascendc.h" + +constexpr int sizeFP16 = 2; +constexpr int alignSizeB = 32; // 32B UB / DataCopy alignment unit +// ALIGN_NUM mirrors the kernel constant (16 FP16 / 32B). Kept local here so the +// host does not need the kernel's constexpr in scope. +constexpr int32_t kAlignNum = alignSizeB / sizeFP16; + +namespace optiling { +static ge::graphStatus TilingFunc(gert::TilingContext* context) { + FusedAddRmsNormTilingData tiling; + + // --- Read shapes: x is (B, H); weight is (H,). --- + const gert::StorageShape* x_shape = context->GetInputShape(0); + if (x_shape == nullptr) { return ge::GRAPH_FAILED; } + const auto& xs = x_shape->GetStorageShape(); + int32_t numDims = static_cast(xs.GetDimNum()); + + int32_t B = 1; + int32_t H = 1; + if (numDims == 0) { + // scalar input — treat as a single element (degenerate, but safe). + B = 1; H = 1; + } else if (numDims == 1) { + // (H,) — a single row. + B = 1; + H = static_cast(xs.GetDim(0)); + } else { + // (B, H) (or higher rank — collapse the leading axes into B). + H = static_cast(xs.GetDim(numDims - 1)); + B = 1; + for (int32_t i = 0; i < numDims - 1; ++i) { + B *= static_cast(xs.GetDim(i)); + } + } + if (H <= 0) H = 1; + if (B <= 0) B = 1; + + // 32B-aligned H (in FP16 elements). UB tiles and vector op counts use this; + // the reduce runs on the exact H (tail ignored). + int32_t alignedHidden = (H + kAlignNum - 1) / kAlignNum * kAlignNum; + + // --- Read eps (OPTIONAL, default 1e-6). enable_pdl is unused by the kernel. --- + float eps = 1e-6f; + const gert::RuntimeAttrs* attrs = context->GetAttrs(); + if (attrs != nullptr) { + const float* epsPtr = attrs->GetFloat(0); // eps (Float, index 0) + if (epsPtr != nullptr) eps = *epsPtr; + } + + tiling.set_batchSize(B); + tiling.set_hiddenSize(H); + tiling.set_alignedHidden(alignedHidden); + tiling.set_alignNum(kAlignNum); + tiling.set_eps(eps); + + // --- Multi-core: split rows across all AIV cores. --- + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); + context->SetBlockDim(ascendcPlatform.GetCoreNumAiv()); + + tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), + context->GetRawTilingData()->GetCapacity()); + context->GetRawTilingData()->SetDataSize(tiling.GetDataSize()); + return ge::GRAPH_SUCCESS; +} +} + + +namespace ge { +// Output shapes mirror the x/residual input (B, H). y and residual_out both have +// the same shape as x. +static ge::graphStatus InferShape(gert::InferShapeContext* context) { + const gert::Shape* x_shape = context->GetInputShape(0); + if (x_shape == nullptr) { return GRAPH_FAILED; } + gert::Shape* y_shape = context->GetOutputShape(0); + gert::Shape* resout_shape = context->GetOutputShape(1); + if (y_shape == nullptr || resout_shape == nullptr) { return GRAPH_FAILED; } + + *y_shape = *x_shape; + *resout_shape = *x_shape; + return GRAPH_SUCCESS; +} + +static ge::graphStatus InferDataType(gert::InferDataTypeContext* context) { + // y and residual_out share x's dtype (FP16). + context->SetOutputDataType(0, context->GetInputDataType(0)); + context->SetOutputDataType(1, context->GetInputDataType(0)); + return ge::GRAPH_SUCCESS; +} +} + + +namespace ops { +class FusedAddRmsNorm : public OpDef { +public: + explicit FusedAddRmsNorm(const char* name) : OpDef(name) + { + this->Input("x") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Input("residual") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Input("weight") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Output("y") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Output("residual_out") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Attr("eps").AttrType(OPTIONAL).Float(1e-06); + this->Attr("enable_pdl").AttrType(OPTIONAL).Bool(false); + + this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType); + + this->AICore() + .SetTiling(optiling::TilingFunc); + this->AICore().AddConfig("ascend910b"); + } +}; + +OP_ADD(FusedAddRmsNorm); +} diff --git a/src/lab3p5/src/ascendc/op_kernel/CMakeLists.txt b/src/lab3p5/src/ascendc/op_kernel/CMakeLists.txt new file mode 100644 index 00000000..609362ee --- /dev/null +++ b/src/lab3p5/src/ascendc/op_kernel/CMakeLists.txt @@ -0,0 +1,23 @@ +# ---------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ---------------------------------------------------------------------------------------------------------- + + +npu_op_kernel_sources(ascendc_kernels + KERNEL_DIR ./ +) + +npu_op_kernel_library(ascendc_kernels + SRC_BASE ${CMAKE_CURRENT_SOURCE_DIR}/ + TILING_LIBRARY cust_optiling +) + +npu_op_package_add(${package_name} + LIBRARY ascendc_kernels +) diff --git a/src/lab3p5/src/ascendc/op_kernel/fused_add_rms_norm.cpp b/src/lab3p5/src/ascendc/op_kernel/fused_add_rms_norm.cpp new file mode 100644 index 00000000..832c3c40 --- /dev/null +++ b/src/lab3p5/src/ascendc/op_kernel/fused_add_rms_norm.cpp @@ -0,0 +1,475 @@ +/** + * @file fused_add_rms_norm.cpp + * @brief Kernel implementation of FusedAddRmsNorm on Ascend C (910B). + * @details + * Op: FusedAddRmsNorm(x, residual, weight, eps) -> (y, residual_out) + * residual_out = x + residual + * y = residual_out / sqrt(mean(residual_out^2, dim=-1) + eps) * weight + * + * Row-parallel baseline (correctness-first, not perf-tuned): + * - Each AIV core owns a contiguous range of rows (block-strided over B). + * - Per row: load x and residual (FP16 GM -> UB via DataCopyPad, zero-padded + * to the 32B-aligned H), cast to FP32, add -> residual_out (FP32); cast + * residual_out back to FP16 and store to GM; square + reduce (FP32 + * BlockReduceSum/WholeReduceSum) to one scalar sumSq; rstd = + * rsqrt(meanSq + eps) computed via the vector Rsqrt on a broadcast tile; + * y = residual_out * rstd * weight, cast to FP16, store to GM. + * - All arithmetic in FP32 for precision; FP16 only on the GM boundary. + * + * Tail handling: H need not be a multiple of 16. UB tiles are sized to + * alignedHidden (rounded up to ALIGN_NUM); DataCopyPad zero-pads the read so + * the tail is well-defined, and the reduce runs on exactly hiddenSize (the + * padded tail is excluded via the SetMaskCount/SetVectorMask path + * used by ReduceNormal). Vector ops run on alignedHidden; the padded tail + * carries garbage but is never written back (CopyOut uses blockLen = H bytes). + * + * Row size: when alignedHidden fits in one UB tile (the lab cases, H<=4096), + * a whole-row path is taken. Larger rows stream through the tile in chunks + * (two passes: pass 1 accumulates sumSq, pass 2 applies rstd*weight and + * writes both outputs), with the weight streamed per chunk. + */ +#include "kernel_operator.h" + +namespace { +constexpr int32_t BUFFER_NUM = 2; // double-buffered in/out queues +// 32B / sizeof(half) == 16: the UB / DataCopy / vector-op alignment unit. +constexpr int32_t ALIGN_NUM = 16; +// UB tile cap (FP32 elements). 910B4 UB = 192 KiB; one FP32 tile of 16 KiB +// (4096 elems) is small and well within budget even with double buffering. +constexpr int32_t TILE_MAX_ELEMS = 4096; +} + +/** + * @brief FusedAddRmsNorm kernel class (FP16 I/O, FP32 compute, row-parallel). + */ +class KernelFusedAddRmsNorm { +public: + __aicore__ inline KernelFusedAddRmsNorm() {} + + __aicore__ inline void Init(GM_ADDR x, GM_ADDR residual, GM_ADDR weight, + GM_ADDR y, GM_ADDR residual_out, + FusedAddRmsNormTilingData& tiling, AscendC::TPipe* pipeIn) { + this->pipe = pipeIn; + this->blockIdx = AscendC::GetBlockIdx(); + + this->batchSize = tiling.batchSize; + this->hiddenSize = tiling.hiddenSize; + this->alignedHidden = tiling.alignedHidden; + this->alignNum = tiling.alignNum; + this->eps = tiling.eps; + + // Per-row UB footprint (capped so a single row tile fits in UB even when + // H is huge; rows larger than this stream in chunks). alignedHidden is a + // multiple of ALIGN_NUM, and TILE_MAX_ELEMS is 4096 == 16*256, so + // tileElems is always a multiple of ALIGN_NUM. + this->tileElems = this->alignedHidden; + if (this->tileElems > TILE_MAX_ELEMS) this->tileElems = TILE_MAX_ELEMS; + if (this->tileElems < this->alignNum) this->tileElems = this->alignNum; + + // Row-parallel split: contiguous range of rows per core. + int32_t totalRows = this->batchSize; + int32_t blockNum = static_cast(AscendC::GetBlockNum()); + int32_t rowsPerBlock = (totalRows + blockNum - 1) / blockNum; + this->startRow = static_cast(this->blockIdx) * rowsPerBlock; + this->endRow = this->startRow + rowsPerBlock; + if (this->endRow > totalRows) this->endRow = totalRows; + + // GM tensors (element counts guarded against 0). + uint64_t totalElems = static_cast(this->batchSize) * + static_cast(this->hiddenSize); + if (totalElems == 0) totalElems = 1; + xGm.SetGlobalBuffer(reinterpret_cast<__gm__ half*>(x), totalElems); + residualGm.SetGlobalBuffer(reinterpret_cast<__gm__ half*>(residual), totalElems); + yGm.SetGlobalBuffer(reinterpret_cast<__gm__ half*>(y), totalElems); + residualOutGm.SetGlobalBuffer(reinterpret_cast<__gm__ half*>(residual_out), totalElems); + uint64_t weightElems = static_cast(this->hiddenSize > 0 ? this->hiddenSize : 1); + weightGm.SetGlobalBuffer(reinterpret_cast<__gm__ half*>(weight), weightElems); + + // UB buffers. + // inQueX / inQueRes : FP16 input tiles (double-buffered). + // outQueY / outQueResOut : FP16 output tiles (double-buffered). + // weightHalfBuf : weight row in FP16 (loaded per chunk / once). + // weightFp32Buf : weight in FP32 (Mul operand). + // resoFp32Buf : residual_out in FP32 (add result; reused as y source). + // sqBuf : squared FP32 tile (reduce source) / rstd broadcast. + // scalarBuf : one FP32 scalar (reduce result staging). + // reduceTmpBuf : scratch required by the Block/Whole reduce intrinsics. + uint32_t tileBytesFp16 = static_cast(this->tileElems) * sizeof(half); + uint32_t tileBytesFp32 = static_cast(this->tileElems) * sizeof(float); + pipe->InitBuffer(inQueX, BUFFER_NUM, tileBytesFp16); + pipe->InitBuffer(inQueRes, BUFFER_NUM, tileBytesFp16); + pipe->InitBuffer(outQueY, BUFFER_NUM, tileBytesFp16); + pipe->InitBuffer(outQueResOut, BUFFER_NUM, tileBytesFp16); + pipe->InitBuffer(weightHalfBuf, tileBytesFp16); + pipe->InitBuffer(weightFp32Buf, tileBytesFp32); + pipe->InitBuffer(resoFp32Buf, tileBytesFp32); + pipe->InitBuffer(sqBuf, tileBytesFp32); + pipe->InitBuffer(scalarBuf, 32); // 1 FP32 scalar, 32B-aligned + pipe->InitBuffer(reduceTmpBuf, 32); // reduce scratch, 32B-aligned + } + + __aicore__ inline void Process() { + if (this->startRow >= this->endRow) return; + if (this->hiddenSize <= 0) return; + + if (this->alignedHidden <= this->tileElems) { + ProcessWholeRows(); + } else { + ProcessChunkedRows(); + } + } + +private: + // ------------------------------------------------------------------ + // Whole-row path (H fits in one UB tile) + // ------------------------------------------------------------------ + __aicore__ inline void ProcessWholeRows() { + AscendC::LocalTensor weightFp32 = weightFp32Buf.Get(); + AscendC::LocalTensor resoFp32 = resoFp32Buf.Get(); + AscendC::LocalTensor sq = sqBuf.Get(); + AscendC::LocalTensor scalar = scalarBuf.Get(); + + // Weight loaded once (full row, zero-padded to alignH), reused per row. + LoadWeightRow(weightFp32, 0, this->hiddenSize, this->alignedHidden); + + const int32_t H = this->hiddenSize; + const int32_t alignH = this->alignedHidden; + const float invH = 1.0f / static_cast(H); + + for (int64_t row = this->startRow; row < this->endRow; ++row) { + uint64_t base = static_cast(row) * static_cast(H); + + // --- Load x, residual (FP16 GM -> UB) --- + AscendC::LocalTensor xLocal = inQueX.AllocTensor(); + AscendC::LocalTensor resLocal = inQueRes.AllocTensor(); + CopyInRow(xLocal, xGm, base); + CopyInRow(resLocal, residualGm, base); + inQueX.EnQue(xLocal); + inQueRes.EnQue(resLocal); + xLocal = inQueX.DeQue(); + resLocal = inQueRes.DeQue(); + + // residual_out (FP32) = Cast(x) + Cast(residual) + AscendC::Cast(resoFp32, xLocal, AscendC::RoundMode::CAST_NONE, alignH); + AscendC::PipeBarrier(); + AscendC::Cast(sq, resLocal, AscendC::RoundMode::CAST_NONE, alignH); + AscendC::PipeBarrier(); + AscendC::Add(resoFp32, resoFp32, sq, alignH); + AscendC::PipeBarrier(); + inQueX.FreeTensor(xLocal); + inQueRes.FreeTensor(resLocal); + + // --- Write residual_out (FP32 -> FP16 GM) --- + AscendC::LocalTensor resOutLocal = outQueResOut.AllocTensor(); + AscendC::Cast(resOutLocal, resoFp32, AscendC::RoundMode::CAST_NONE, alignH); + AscendC::PipeBarrier(); + outQueResOut.EnQue(resOutLocal); + resOutLocal = outQueResOut.DeQue(); + CopyOutRow(resOutLocal, residualOutGm, base); + outQueResOut.FreeTensor(resOutLocal); + + // --- Reduce sum(residual_out^2) over the row (FP32) --- + AscendC::Mul(sq, resoFp32, resoFp32, alignH); + AscendC::PipeBarrier(); + ReduceNormal(scalar, sq, H); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + float sumSq = scalar.GetValue(0); + + // rstd = 1 / sqrt(meanSq + eps). The 910B Rsqrt intrinsic is a fast + // approximation (~0.2% error, enough to blow the fp16 tol); instead + // compute rms = Sqrt(mean+eps) and divide, matching the fp32 golden + // (torch.sqrt then divide) closely. Sqrt+Div are both ~1-ULP. + float meanPlusEps = sumSq * invH + this->eps; + AscendC::Duplicate(sq, meanPlusEps, alignH); + AscendC::PipeBarrier(); + AscendC::Sqrt(sq, sq, alignH); // sq = rms + AscendC::PipeBarrier(); + + // --- y = (residual_out / rms) * weight (FP32), cast FP16, write GM --- + AscendC::Div(resoFp32, resoFp32, sq, alignH); // /= rms + AscendC::PipeBarrier(); + AscendC::Mul(resoFp32, resoFp32, weightFp32, alignH); // *= weight + AscendC::PipeBarrier(); + AscendC::LocalTensor yLocal = outQueY.AllocTensor(); + AscendC::Cast(yLocal, resoFp32, AscendC::RoundMode::CAST_NONE, alignH); + AscendC::PipeBarrier(); + outQueY.EnQue(yLocal); + yLocal = outQueY.DeQue(); + CopyOutRow(yLocal, yGm, base); + outQueY.FreeTensor(yLocal); + } + } + + // ------------------------------------------------------------------ + // Chunked-row path (H > UB tile): two streaming passes per row. + // Pass 1: stream chunks, accumulate sum(residual_out^2) -> rstd. + // Pass 2: stream chunks, apply rstd*weight, write y + residual_out. + // (Weight is streamed per chunk in pass 2.) + // ------------------------------------------------------------------ + __aicore__ inline void ProcessChunkedRows() { + AscendC::LocalTensor weightFp32 = weightFp32Buf.Get(); + AscendC::LocalTensor resoFp32 = resoFp32Buf.Get(); + AscendC::LocalTensor sq = sqBuf.Get(); + AscendC::LocalTensor scalar = scalarBuf.Get(); + + const int32_t H = this->hiddenSize; + const int32_t chunkElems = this->tileElems; // multiple of ALIGN_NUM + const float invH = 1.0f / static_cast(H); + + for (int64_t row = this->startRow; row < this->endRow; ++row) { + uint64_t base = static_cast(row) * static_cast(H); + + // --- Pass 1: residual_out + accumulate sum-of-squares --- + float sumSq = 0.0f; + int32_t off = 0; + while (off < H) { + int32_t n = (H - off > chunkElems) ? chunkElems : (H - off); + int32_t nAlign = (n + this->alignNum - 1) / this->alignNum * this->alignNum; + + AscendC::LocalTensor xLocal = inQueX.AllocTensor(); + AscendC::LocalTensor resLocal = inQueRes.AllocTensor(); + CopyInChunk(xLocal, xGm, base + off, n, nAlign); + CopyInChunk(resLocal, residualGm, base + off, n, nAlign); + inQueX.EnQue(xLocal); + inQueRes.EnQue(resLocal); + xLocal = inQueX.DeQue(); + resLocal = inQueRes.DeQue(); + + AscendC::Cast(resoFp32, xLocal, AscendC::RoundMode::CAST_NONE, nAlign); + AscendC::PipeBarrier(); + AscendC::Cast(sq, resLocal, AscendC::RoundMode::CAST_NONE, nAlign); + AscendC::PipeBarrier(); + AscendC::Add(resoFp32, resoFp32, sq, nAlign); + AscendC::PipeBarrier(); + AscendC::Mul(sq, resoFp32, resoFp32, nAlign); + AscendC::PipeBarrier(); + ReduceNormal(scalar, sq, n); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + sumSq += scalar.GetValue(0); + + inQueX.FreeTensor(xLocal); + inQueRes.FreeTensor(resLocal); + off += n; + } + + float meanPlusEps = sumSq * invH + this->eps; + + // --- Pass 2: recompute residual_out, apply rstd*weight, write y + res_out --- + off = 0; + while (off < H) { + int32_t n = (H - off > chunkElems) ? chunkElems : (H - off); + int32_t nAlign = (n + this->alignNum - 1) / this->alignNum * this->alignNum; + + AscendC::LocalTensor xLocal = inQueX.AllocTensor(); + AscendC::LocalTensor resLocal = inQueRes.AllocTensor(); + CopyInChunk(xLocal, xGm, base + off, n, nAlign); + CopyInChunk(resLocal, residualGm, base + off, n, nAlign); + inQueX.EnQue(xLocal); + inQueRes.EnQue(resLocal); + xLocal = inQueX.DeQue(); + resLocal = inQueRes.DeQue(); + + AscendC::Cast(resoFp32, xLocal, AscendC::RoundMode::CAST_NONE, nAlign); + AscendC::PipeBarrier(); + AscendC::Cast(sq, resLocal, AscendC::RoundMode::CAST_NONE, nAlign); + AscendC::PipeBarrier(); + AscendC::Add(resoFp32, resoFp32, sq, nAlign); + AscendC::PipeBarrier(); + inQueX.FreeTensor(xLocal); + inQueRes.FreeTensor(resLocal); + + // residual_out -> GM (FP16) + AscendC::LocalTensor resOutLocal = outQueResOut.AllocTensor(); + AscendC::Cast(resOutLocal, resoFp32, AscendC::RoundMode::CAST_NONE, nAlign); + AscendC::PipeBarrier(); + outQueResOut.EnQue(resOutLocal); + resOutLocal = outQueResOut.DeQue(); + CopyOutChunk(resOutLocal, residualOutGm, base + off, n); + outQueResOut.FreeTensor(resOutLocal); + + // y = (residual_out / rms) * weight (Div+Sqrt path, see whole-row). + AscendC::Duplicate(sq, meanPlusEps, nAlign); + AscendC::PipeBarrier(); + AscendC::Sqrt(sq, sq, nAlign); + AscendC::PipeBarrier(); + AscendC::Div(resoFp32, resoFp32, sq, nAlign); + AscendC::PipeBarrier(); + LoadWeightRow(weightFp32, off, n, nAlign); // weight chunk for this offset + AscendC::Mul(resoFp32, resoFp32, weightFp32, nAlign); + AscendC::PipeBarrier(); + + AscendC::LocalTensor yLocal = outQueY.AllocTensor(); + AscendC::Cast(yLocal, resoFp32, AscendC::RoundMode::CAST_NONE, nAlign); + AscendC::PipeBarrier(); + outQueY.EnQue(yLocal); + yLocal = outQueY.DeQue(); + CopyOutChunk(yLocal, yGm, base + off, n); + outQueY.FreeTensor(yLocal); + + off += n; + } + } + } + + // ------------------------------------------------------------------ + // Helpers + // ------------------------------------------------------------------ + // Load `realN` weight elements from GM offset `off`, zero-pad the tail up + // to `nAlign` (nAlign >= realN, both multiples of ALIGN_NUM), and Cast them + // into the FP32 weight tile `wFp32[0..nAlign)`. The padded tail is zero, so + // the later Mul(reso, reso, wFp32, nAlign) is correct for the real elements + // and harmless (×0) for the tail — which is never written back anyway + // (CopyOut uses blockLen = n bytes). + __aicore__ inline void LoadWeightRow(AscendC::LocalTensor& wFp32, + int32_t off, int32_t realN, int32_t nAlign) { + AscendC::LocalTensor wHalf = weightHalfBuf.Get(); + AscendC::DataCopyExtParams copyParams; + copyParams.blockCount = 1; + copyParams.blockLen = static_cast(realN * sizeof(half)); + copyParams.srcStride = 0; + copyParams.dstStride = 0; + AscendC::DataCopyPadExtParams padParams; + padParams.isPad = (realN < nAlign); + padParams.leftPadding = 0; + padParams.rightPadding = static_cast(nAlign - realN); + padParams.paddingValue = 0; + AscendC::DataCopyPad(wHalf, weightGm[static_cast(off)], copyParams, padParams); + AscendC::PipeBarrier(); + AscendC::Cast(wFp32, wHalf, AscendC::RoundMode::CAST_NONE, nAlign); + AscendC::PipeBarrier(); + } + + // Copy a full aligned row (alignedHidden elems) from GM half -> UB half, + // with zero-padding of the tail when hiddenSize < alignedHidden. + __aicore__ inline void CopyInRow(AscendC::LocalTensor& dst, + AscendC::GlobalTensor& src, uint64_t off) { + AscendC::DataCopyExtParams copyParams; + copyParams.blockCount = 1; + copyParams.blockLen = static_cast(this->hiddenSize * sizeof(half)); + copyParams.srcStride = 0; + copyParams.dstStride = 0; + AscendC::DataCopyPadExtParams padParams; + padParams.isPad = (this->hiddenSize < this->alignedHidden); + padParams.leftPadding = 0; + padParams.rightPadding = static_cast(this->alignedHidden - this->hiddenSize); + padParams.paddingValue = 0; + AscendC::DataCopyPad(dst, src[off], copyParams, padParams); + AscendC::PipeBarrier(); + } + + // Copy `n` elems (padded to nAlign) from GM half -> UB half. + __aicore__ inline void CopyInChunk(AscendC::LocalTensor& dst, + AscendC::GlobalTensor& src, + uint64_t off, int32_t n, int32_t nAlign) { + AscendC::DataCopyExtParams copyParams; + copyParams.blockCount = 1; + copyParams.blockLen = static_cast(n * sizeof(half)); + copyParams.srcStride = 0; + copyParams.dstStride = 0; + AscendC::DataCopyPadExtParams padParams; + padParams.isPad = (n < nAlign); + padParams.leftPadding = 0; + padParams.rightPadding = static_cast(nAlign - n); + padParams.paddingValue = 0; + AscendC::DataCopyPad(dst, src[off], copyParams, padParams); + AscendC::PipeBarrier(); + } + + // Copy a full row (hiddenSize elems) from UB half -> GM half. Only the first + // hiddenSize elements are written (blockLen = hiddenSize * sizeof(half) bytes). + __aicore__ inline void CopyOutRow(AscendC::LocalTensor& src, + AscendC::GlobalTensor& dst, uint64_t off) { + AscendC::DataCopyExtParams copyParams; + copyParams.blockCount = 1; + copyParams.blockLen = static_cast(this->hiddenSize * sizeof(half)); + copyParams.srcStride = 0; + copyParams.dstStride = 0; + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::DataCopyPad(dst[off], src, copyParams); + AscendC::PipeBarrier(); + } + + // Copy `n` elems from UB half -> GM half (byte-granular blockLen). + __aicore__ inline void CopyOutChunk(AscendC::LocalTensor& src, + AscendC::GlobalTensor& dst, + uint64_t off, int32_t n) { + AscendC::DataCopyExtParams copyParams; + copyParams.blockCount = 1; + copyParams.blockLen = static_cast(n * sizeof(half)); + copyParams.srcStride = 0; + copyParams.dstStride = 0; + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::DataCopyPad(dst[off], src, copyParams); + AscendC::PipeBarrier(); + } + + // High-performance FP32 reduce: BlockReduceSum loop + WholeReduceSum. + // Reduces the first `totalElements` of src into dst[0] (one FP32 scalar). + // Uses SetMaskCount + SetVectorMask(totalElements) so only the + // first totalElements participate (the UB tail, if any, is ignored). + // NOTE: this clobbers `src` in place (BlockReduceSum writes partial sums + // back into it); callers must not rely on src afterwards. + __aicore__ inline void ReduceNormal(const AscendC::LocalTensor& dst, + const AscendC::LocalTensor& src, + const int totalElements) { + constexpr int elemsPerBlock = 32 / sizeof(float); // 8 + int currentLen = totalElements; + AscendC::SetMaskCount(); + while (currentLen > (elemsPerBlock * 8)) { + int blockCount = (currentLen + elemsPerBlock - 1) / elemsPerBlock; + int repeat = (blockCount + 7) / 8; + AscendC::SetVectorMask(currentLen); + AscendC::BlockReduceSum(src, src, repeat, + AscendC::MASK_PLACEHOLDER, 1, 1, 8); + currentLen = blockCount; + } + AscendC::SetVectorMask(currentLen); + AscendC::WholeReduceSum(dst, src, AscendC::MASK_PLACEHOLDER, 1, 1, 1, 8); + AscendC::SetMaskNorm(); + AscendC::ResetMask(); + } + +private: + AscendC::TPipe* pipe; + int32_t blockIdx; + int32_t batchSize; + int32_t hiddenSize; + int32_t alignedHidden; + int32_t alignNum; + int32_t tileElems; + int64_t startRow; + int64_t endRow; + float eps; + + AscendC::GlobalTensor xGm; + AscendC::GlobalTensor residualGm; + AscendC::GlobalTensor weightGm; + AscendC::GlobalTensor yGm; + AscendC::GlobalTensor residualOutGm; + + AscendC::TQue inQueX; + AscendC::TQue inQueRes; + AscendC::TQue outQueY; + AscendC::TQue outQueResOut; + AscendC::TBuf weightHalfBuf; + AscendC::TBuf weightFp32Buf; + AscendC::TBuf resoFp32Buf; + AscendC::TBuf sqBuf; + AscendC::TBuf scalarBuf; + AscendC::TBuf reduceTmpBuf; +}; + + +extern "C" __global__ __aicore__ void fused_add_rms_norm(GM_ADDR x, GM_ADDR residual, GM_ADDR weight, + GM_ADDR y, GM_ADDR residual_out, + GM_ADDR workspace, GM_ADDR tiling) { + GET_TILING_DATA(tilingData, tiling); + AscendC::TPipe pipe; + KernelFusedAddRmsNorm op; + op.Init(x, residual, weight, y, residual_out, tilingData, &pipe); + op.Process(); +} diff --git a/src/lab3p5/src/ascendc/op_kernel/fused_add_rms_norm_tiling.h b/src/lab3p5/src/ascendc/op_kernel/fused_add_rms_norm_tiling.h new file mode 100644 index 00000000..62c01380 --- /dev/null +++ b/src/lab3p5/src/ascendc/op_kernel/fused_add_rms_norm_tiling.h @@ -0,0 +1,43 @@ +/** + * @file fused_add_rms_norm_tiling.h + * @brief Tiling data definition for the FusedAddRmsNorm operator + * + * Op: FusedAddRmsNorm(x, residual, weight, eps, enable_pdl) -> (y, residual_out) + * residual_out = x + residual + * y = residual_out / sqrt(mean(residual_out^2, dim=-1) + eps) * weight + * + * Inputs/outputs are FP16 ND tensors of shape (B, H); weight is (H,). The kernel + * is row-parallel (one row per work item, split across AIV cores) and computes + * in FP32 for precision, casting back to FP16 on the way out. + * + * Only op_host includes this header. The kernel accesses the tiling fields + * through the GET_TILING_DATA(tilingData, tiling) macro (it must NOT include + * this file — that would pull in graph/types.h and break the kernel compile). + */ +#ifndef FUSED_ADD_RMS_NORM_TILING_H +#define FUSED_ADD_RMS_NORM_TILING_H + +#include "register/tilingdata_base.h" + +namespace optiling { +// 32B UB alignment unit == 16 FP16 elements == 8 FP32 elements. +constexpr int32_t ALIGN_NUM = 16; + +BEGIN_TILING_DATA_DEF(FusedAddRmsNormTilingData) + // Number of rows (B) in the (B, H) input. + TILING_DATA_FIELD_DEF(int32_t, batchSize); + // Hidden size (H) — the per-row reduction length. + TILING_DATA_FIELD_DEF(int32_t, hiddenSize); + // hiddenSize rounded up to a multiple of ALIGN_NUM (FP16 32B alignment). + // UB tiles are sized to this; vector ops run on alignedHidden, reduce runs + // on hiddenSize (the tail is masked out / ignored). + TILING_DATA_FIELD_DEF(int32_t, alignedHidden); + // 32B alignment unit in elements (== ALIGN_NUM, mirrored for the kernel). + TILING_DATA_FIELD_DEF(int32_t, alignNum); + // eps added inside the sqrt for numerical stability. + TILING_DATA_FIELD_DEF(float, eps); +END_TILING_DATA_DEF; + +REGISTER_TILING_DATA_CLASS(FusedAddRmsNorm, FusedAddRmsNormTilingData) +} +#endif // FUSED_ADD_RMS_NORM_TILING_H diff --git a/src/lab3p5/src/ascendc/setup.py b/src/lab3p5/src/ascendc/setup.py new file mode 100644 index 00000000..abc4ffcd --- /dev/null +++ b/src/lab3p5/src/ascendc/setup.py @@ -0,0 +1,35 @@ +"""Build the student's custom-op extension into a wheel and install it. + +Mirrors the S9 `Transpose/setup.py` exactly. Produces a `custom_ops_lib` +module that test_op.py imports. +""" +import os +import torch +from setuptools import setup, find_packages +from torch.utils.cpp_extension import BuildExtension + +import torch_npu +from torch_npu.utils.cpp_extension import NpuExtension + +PYTORCH_NPU_INSTALL_PATH = os.path.dirname(os.path.abspath(torch_npu.__file__)) + +exts = [] +ext1 = NpuExtension( + name="custom_ops_lib", + sources=["./extension/custom_op.cpp"], + extra_compile_args=[ + '-I' + os.path.join(PYTORCH_NPU_INSTALL_PATH, "include/third_party/acl/inc"), + # Make `../common/pytorch_npu_helper.hpp` resolvable from extension/custom_op.cpp. + '-I' + os.path.dirname(os.path.abspath(__file__)), + ], +) +exts.append(ext1) + +setup( + name="custom_ops", + version='1.0', + keywords='custom_ops', + ext_modules=exts, + packages=find_packages(), + cmdclass={"build_ext": BuildExtension}, +) diff --git a/src/lab3p5/src/tilelang/__init__.py b/src/lab3p5/src/tilelang/__init__.py new file mode 100644 index 00000000..b113492e --- /dev/null +++ b/src/lab3p5/src/tilelang/__init__.py @@ -0,0 +1,15 @@ +"""TileLang-Ascend entry. When USE_TILELANG=1, test_op.py imports this so that +`custom_ops_lib.fused_add_rmsnorm(...)` resolves to the TileLang kernel below. +Requires tilelang-ascend (Ascend target) installed. +""" +import sys + +from .fused_add_rmsnorm import fused_add_rmsnorm, BackendUnavailable + + +class _Lib: + def fused_add_rmsnorm(self, x, residual, weight, eps): + return fused_add_rmsnorm(x, residual, weight, eps) + + +sys.modules["custom_ops_lib"] = _Lib() diff --git a/src/lab3p5/src/tilelang/fused_add_rmsnorm.py b/src/lab3p5/src/tilelang/fused_add_rmsnorm.py new file mode 100644 index 00000000..22f43329 --- /dev/null +++ b/src/lab3p5/src/tilelang/fused_add_rmsnorm.py @@ -0,0 +1,51 @@ +"""TileLang-Ascend FusedAddRmsNorm kernel — skeleton. + +Students: implement `_build_kernel` (the `@T.prim_func` body) and the +`fused_add_rmsnorm` launcher. The tilelang-ascend backend must be installed +(separate `tilelang-ascend` package) — otherwise `BackendUnavailable` is +raised and the harness skips this backend. +""" +from __future__ import annotations + + +class BackendUnavailable(RuntimeError): + """Raised when the TileLang Ascend backend is not installed.""" + + +_backend_ok: bool | None = None +_kernel_cache: dict = {} + +try: + import tilelang + import tilelang.language as T +except Exception: # pragma: no cover + tilelang = None + T = None + + +def _ascend_target_available() -> bool: + """True iff the ascend target detector is registered (tilelang-ascend installed).""" + if tilelang is None: + return False + try: + from tilelang.backend.target import list_target_detectors + return any("ascend" in str(d).lower() for d in list_target_detectors()) + except Exception: + return False + + +def _check_backend() -> bool: + global _backend_ok + if _backend_ok is not None: + return _backend_ok + _backend_ok = _ascend_target_available() + return _backend_ok + + +def fused_add_rmsnorm(x, residual, weight, eps: float = 1e-6): + """x,residual: (B,H) fp16; weight: (H,) fp16. Returns (y, residual_out).""" + # TODO: implement + raise NotImplementedError("fused_add_rmsnorm not implemented") + + +KERNEL_NAME_HINT = "addrmsnorm" diff --git a/src/lab3p5/src/triton/__init__.py b/src/lab3p5/src/triton/__init__.py new file mode 100644 index 00000000..9c14173f --- /dev/null +++ b/src/lab3p5/src/triton/__init__.py @@ -0,0 +1,15 @@ +"""Triton-Ascend entry. When USE_TRITON=1, test_op.py imports this so that +`custom_ops_lib.fused_add_rmsnorm(...)` resolves to the Triton kernel below — +no wheel build needed. +""" +import sys + +from .fused_add_rmsnorm import fused_add_rmsnorm + + +class _Lib: + def fused_add_rmsnorm(self, x, residual, weight, eps): + return fused_add_rmsnorm(x, residual, weight, eps) + + +sys.modules["custom_ops_lib"] = _Lib() diff --git a/src/lab3p5/src/triton/fused_add_rmsnorm.py b/src/lab3p5/src/triton/fused_add_rmsnorm.py new file mode 100644 index 00000000..3003726e --- /dev/null +++ b/src/lab3p5/src/triton/fused_add_rmsnorm.py @@ -0,0 +1,32 @@ +"""FusedAddRmsNorm — Triton-Ascend kernel skeleton. + +Students: implement `_fused_add_rmsnorm_kernel` and the launch in +`fused_add_rmsnorm`. The grid is pinned to the physical vector-core count by +convention; you may adjust it. +""" +import torch +import triton +import triton.language as tl + + +@triton.jit +def _fused_add_rmsnorm_kernel( + x_ptr, residual_ptr, weight_ptr, y_ptr, resout_ptr, + B, H, + eps, + BLOCK_H: tl.constexpr, +): + # TODO: implement kernel + pid = tl.program_id(0) + pass + + +def fused_add_rmsnorm(x, residual, weight, eps: float = 1e-6): + """x,residual: (B,H) fp16 NPU; weight: (H,) fp16. Returns (y, residual_out).""" + # TODO: launch kernel + raise NotImplementedError("fused_add_rmsnorm not implemented") + + +# Expose the same module name the test harness imports. +def custom_op(x, residual, weight, eps): + return fused_add_rmsnorm(x, residual, weight, eps) diff --git a/zensical.toml b/zensical.toml index bb2cd9e2..67a477df 100644 --- a/zensical.toml +++ b/zensical.toml @@ -26,6 +26,7 @@ nav = [ { "Lab 1: 简单集群搭建" = "lab/Lab1-MiniCluster/index.md" }, { "Lab 2: MoE 的向量化计算" = "lab/Lab2-Vectorization/index.md"}, { "Lab 3: GDN Prefill 前向优化" = "lab/Lab3-GDN-Prefill/index.md" }, + { "Lab 3.5: 昇腾算子开发与优化" = "lab/Lab3.5-AscendC-Op/index.md" }, { "Lab 4: AMSS-NCKU 数值相对论程序优化" = "lab/Lab4-AMSS-NCKU/index.md"}, { "Lab 5: Gemma4 端到端推理优化" = "lab/Lab5-Gemma4/index.md"} ] },