算子开发示例
本目录包含 pyasc 算子开发端到端示例,覆盖从基础算子到进阶场景的完整开发流程。所有示例均采用 PyTorch 输入输出 Tensor,可直接运行验证。
示例总览
| 算子样例 | 功能描述 |
|---|---|
| 01_add | 实现手动插入同步流水的Add算子。 |
| 02_add_framework | 实现通过Ascend C框架插入流水同步的Add算子。 |
| 03_matmul_mix | 实现MIX模式(包含矩阵计算和矢量计算)下的Matmul算子,计算公式为:C = A * B。 |
| 04_matmul_cube_only | 实现纯Cube模式(只有矩阵计算)的Matmul算子,计算公式为:C = A * B。 |
| 05_matmul_leakyrelu | 实现MatmulLeakyRelu算子,计算公式为:C = A * B + Bias, C = C > 0 ? C : C * 0.001。 |
| 06_gelu | 实现GELU激活函数算子,计算公式为:GELU(x) = x * Φ(x),采用tanh近似。 |
| 07_swiglu | 实现 SwiGLU 激活函数算子,TPipe/TQue 框架双路 VECIN 输入,五步基础向量算子组合完成计算。 |
| 08_rmsnorm | 实现 RMSNorm 归一化算子,TPipe/TQue 框架管理搬运,调用 asc.adv.rmsnorm 高阶 API 完成行归一化计算。 |
推荐学习顺序
建议按以下顺序循序渐进地学习:
- 01_add — 入门首选。理解手动同步流水、数据搬运(Global↔Local)、多核切分(tiling)的基本概念。
- 02_add_framework — 学习 Ascend C 框架提供的 TPipe/TQue/TBuf 机制,体会自动流水同步如何简化开发。
- 03_matmul_mix — 进入矩阵计算领域,掌握 Matmul 高阶 API 的 MIX 模式用法。
- 04_matmul_cube_only — 深入理解纯 Cube 模式下的矩阵分块与 Workspace 配置。
- 05_matmul_leakyrelu — 学习融合算子开发,将矩阵乘与激活函数在单次 Kernel 调用中完成,避免中间数据回传。
- 06_gelu — 学习复杂激活函数的实现,掌握九步算子组合、TBuf 复用和 adv API(tanh)的使用。
- 07_swiglu — 学习 SwiGLU 激活函数:TPipe/TQue 框架同步、双路 VECIN、TBuf 复用、多步组合计算与 PIPE_V 同步。
- 08_rmsnorm — 学习 RMSNorm 归一化:
asc.adv.rmsnorm高阶 API 调用、RmsNormTiling 分块参数配置、整块/尾块处理。
运行示例
环境准备
请先完成环境准备,确保已安装 PyTorch 和 torch_npu 插件。
通用运行命令
cd pyasc/examples/<示例目录>
python3 <示例文件>.py -r [RUN_MODE] -v [SOC_VERSION]
RUN_MODE:Model(仿真器,无需 NPU)或NPU(上板运行)SOC_VERSION:昇腾 AI 处理器型号,如Ascend910B1、Ascend910C
示例
# 仿真器模式运行 Add 示例
python3 examples/01_add/add.py -r Model -v Ascend910B1
# NPU 上板运行 Matmul MIX 示例
python3 examples/03_matmul_mix/matmul_mix.py -r NPU -v Ascend910C