CANN TileLang算子代码生成

【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skillsname: tilelang-op-developer description: 基于设计文档生成 TileLang-Ascend 算子实现代码与测试。从 design.md 中提取关键信息结合 examples/ 中的参考实现生成可运行代码。触发实现算子、写 kernel、生成代码、算子编码、根据设计文档实现。TileLang-Ascend 算子代码生成基于设计文档design.md和已有示例生成可运行的算子实现与测试。1. 从 design.md 中提取的信息只取这些design.md 可能很长只提取以下字段忽略其余内容提取字段所在章节用途数学公式§1 概述理解计算逻辑算法步骤分解§1 算法描述确定计算顺序API 映射表§3 API 映射设计核心每步用哪个 TileLang API伪代码§3 计算伪代码核心代码骨架输入输出 shape 和 dtype§4 数据规格函数签名和测试数据block 大小§5 Tiling 策略分块参数pass_configs§7 同步策略JIT 配置Golden 函数§8 验证方案测试对比基准测试用例表§8 验证方案测试配置精度标准§8 验证方案atol / rtol明确忽略的内容这些容易误导模式选型的分析推理过程内存预算的计算过程和多轮优化迭代风险点与注意事项过于笼统交付清单仅是文件列表任何标注为待确认的内容2. 参考来源优先级高于 design.md 伪代码当 design.md 伪代码与 examples/ 中同类实现有冲突时以 examples/ 为准。2.1 API 用法和模式选择API 用法查阅 tilelang-api-best-practices SKILL.md 及其 references 目录编程模式和 pass_configs查阅 tilelang-programming-model-guide SKILL.md 及其 references 目录2.2 同类算子示例生成代码前必须查阅examples/中的同类算子算子类型参考示例逐元素运算add/mul/sigmoid/reluexamples/elementwise/、examples/activation/归约运算reduce_sum/max/minexamples/reduce/归一化softmax/layernorm/rmsnormexamples/softmax/、examples/normalization/GEMMexamples/gemm/、examples/developer_mode/gemm_developer.py融合算子examples/flash_attention/Developer 模式examples/developer_mode/查阅示例时关注Kernel 结构T.Kernel参数、cid/vid用法Buffer 分配方式shape 和 dtypepass_configs 配置该类算子实际使用哪些开关数据搬运T.copy的索引写法3. 代码生成流程步骤 1读取设计文档读取design.md按 §1 的表格提取字段。步骤 2查找参考示例在examples/中找到最相似的算子实现完整阅读其代码。步骤 3生成实现代码基于 design.md 的 API 映射 参考示例的代码风格生成example_{op}.py。文件结构import tilelang from tilelang import DataType, language as T import torch tilelang.cache.clear_cache() # 算子实现 tilelang.jit(out_idx[...], pass_configs{...}) def op_name(M, N, block_M, block_N, dtypefloat): # 分块计算 m_num T.ceildiv(M, block_M) n_num T.ceildiv(N, block_N) VEC_NUM 2 T.prim_func def main(Input: T.Tensor((M, N), dtype), Output: T.Tensor((M, N), dtype)): with T.Kernel(..., is_npuTrue) as (cid, vid): # buffer 分配 # 数据搬入 # 计算 # 数据搬出 pass return main # 测试 if __name__ __main__: torch.manual_seed(0) test_configs [...] # 来自 design.md §8 for config in test_configs: # 1. 创建 kernel # 2. 生成输入数据 # 3. 执行 kernel # 4. golden 对比 # 5. 精度检查 pass print(All tests passed!)步骤 4运行验证python examples/{op}/example_{op}.py如果报错按以下顺序排查编译错误→ 检查 buffer 大小、API 参数、对齐运行错误→ 检查索引越界、同步缺失精度错误→ 检查计算公式、数据类型、容差设置4. 关键编码规范Buffer 分配# VEC_NUM 2每个 vector 核处理 block_M // VEC_NUM 行 a_ub T.alloc_ub([block_M // VEC_NUM, block_N], dtype)数据搬运索引# 标准索引模式 row_start bx * block_M vid * block_M // VEC_NUM T.copy(A[row_start, by * block_N], a_ub) T.copy(a_ub, B[row_start, by * block_N])同步# Expert 模式手同步 with T.Scope(V): T.copy(A[...], a_ub) T.barrier_all() T.tile.exp(a_ub, a_ub) T.barrier_all() T.copy(a_ub, B[...]) # Developer 模式 自动同步无需手动 barrier pass_configs { tilelang.PassConfigKey.TL_ASCEND_AUTO_SYNC: True, tilelang.PassConfigKey.TL_ASCEND_MEMORY_PLANNING: True, }广播# 归约结果 [M, 1] 广播到 [M, N] max_ub T.alloc_ub([block_M // VEC_NUM, 1], dtype) max_2d_ub T.alloc_ub([block_M // VEC_NUM, block_N], dtype) T.tile.broadcast(max_2d_ub, max_ub)测试模板# golden 对比 ref_output torch.nn.functional.softmax(input_data, dim-1) # 或手写 golden torch.testing.assert_close(output.cpu(), ref_output.cpu(), rtolrtol, atolatol)5. Checklist生成代码后逐项检查#检查项1out_idx与函数签名中的输出参数位置一致2block_M // VEC_NUM在 buffer 分配和索引中一致使用3所有T.alloc_ub的 shape 乘积不超 UB 容量4Expert 模式有T.Scope(V)和T.barrier_all()5Developer 模式有对应的pass_configs6测试包含至少 2 个配置小规模 典型规模7golden 函数使用 PyTorch 标准实现【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考