CANN TileLang算子设计模板

{算子名称} 算子设计文档

【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体，本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills

1. 概述

1.1 算子名称

{算子名称}

1.2 功能描述

{一句话描述算子功能}

1.3 数学公式

$$ {数学公式} $$

1.4 算法描述

{对于多步算子，描述计算步骤的分解逻辑。单步算子可省略。}

1.5 数据流图

输入张量 → [计算步骤1] → [计算步骤2] → ... → 输出张量

2. 编程模式选型

2.1 模式结论

选定模式: {Developer / Expert / 混合}

2.2 选型理由

{基于算子特征（计算类型、是否含 matmul、是否含归约、是否需要流水线）的分析}

2.3 模式影响

3. API 映射设计

3.1 公式拆解

3.2 TileLang API 映射

3.3 计算伪代码

# 基于 TileLang API 的计算流程伪代码 with T.Kernel(block_num, is_npu=True) as (cid, vid): # 1. 分配 buffer {buffer 分配代码} # 2. 数据搬入 {T.copy 搬入代码} # 3. 计算 {核心计算代码} # 4. 数据搬出 {T.copy 搬出代码}

3.4 API 可行性确认

{列出使用的 API 及其来源确认（速查表 / 示例 / 源码），标注是否经过验证}

4. 数据规格与内存规划

4.1 输入张量

4.2 输出张量

4.3 中间缓冲区

4.4 内存搬运路径

{完整的数据搬运路径图} 示例（纯 Vector）: GM[A] --T.copy--> UB[a_ub] --计算--> UB[c_ub] --T.copy--> GM[C] 示例（Cube + Vector）: GM[A] --T.copy--> L1[a_l1] --T.copy--> L0A[a_l0a] GM[B] --T.copy--> L1[b_l1] --T.copy--> L0B[b_l0b] L0A + L0B --T.gemm--> L0C[c_l0c] --T.copy--> UB[c_ub] UB[c_ub] --后处理--> UB[c_ub] --T.copy--> GM[C]

4.5 UB 内存预算

4.6 动态轴定义

{如无动态轴，写"无"。如有：}

4.7 JIT 配置

@tilelang.jit( out_idx=[{输出索引}], pass_configs={ {pass 配置项} }, )

5. Tiling 策略

5.1 计算类型

类型: {纯 Vector / 纯 Cube / 混合}

判定依据: {如: 算子仅包含 element-wise 运算，无 matmul，判定为纯 Vector}

5.2 Block 划分

block_M = {值} # {选择理由} block_N = {值} # {选择理由} block_num = (M // block_M) * (N // block_N)

5.3 约束分析

• 对齐约束: {如: block_N=128, fp16 尾轴 128 > 16 ✓}
• UB 容量: {如: 总 buffer = 64KB < 128KB ✓}
• L0 容量: {如: 无 Cube 计算，不适用}

5.4 注意事项

{非整除情况的处理策略、边界块的特殊逻辑等}

6. 循环与调度结构

6.1 循环结构总结

6.2 循环伪代码

# Block 级并行（隐式，由 T.Kernel 管理） with T.Kernel(block_num, is_npu=True) as (cid, vid): {block 内循环结构}

6.3 流水线优化

{是否使用 T.Pipelined？如使用，说明 num_stages 设计和 buffer 管理策略}

6.4 尾块处理

{当输入 shape 不能被 block size 整除时的处理策略}

7. 同步策略

7.1 同步模式

模式: {自动同步 / 手动同步 / 混合}

7.2 同步点说明

{手动同步时，列出每个同步点及理由：}

7.3 pass_configs 配置

pass_configs = { {与同步相关的 pass 配置} }

8. 验证方案

8.1 Golden 函数

def golden_{算子名}({参数}): """基于 PyTorch/NumPy 的参考实现""" {参考实现代码}

8.2 测试用例

8.3 精度标准

9. 风险点与注意事项

9.1 已知约束

{列出本算子在 TileLang-Ascend 上的已知限制}

9.2 常见错误

9.3 特殊场景处理

{如: 非整除分块、极小 shape、混合精度等}

10. 交付清单

10.1 目录结构

examples/{算子名}/ ├── example_{算子名}.py # 算子实现 + 简单测试 ├── design.md # 本设计文档 └── README.md # 使用说明（可选）

10.2 文件清单

10.3 命名规范

• 目录名:{算子名}（snake_case）
• 实现文件:example_{算子名}.py
• 测试文件:test_{算子名}.py

10.4 实现顺序

1. ✅ 设计文档（design.md）
2. ⬜ Golden 函数（验证基准）
3. ⬜ 算子实现（example_{算子名}.py）
4. ⬜ 基础测试（Level 0 + Level 1）
5. ⬜ 边界测试（Level 2）
6. ⬜ 性能测试（Level 3，可选）

【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体，本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills