| name | external-cannbot-ops-lab-tilelang-skills-tilelang-op-developer |
| description | 基于设计文档生成 TileLang-Ascend 算子实现代码与测试。从 design.md 中提取关键信息，结合 examples/ 中的参考实现生成可运行代码。触发：实现算子、写 kernel、生成代码、算子编码、根据设计文档实现。 |
| original-name | tilelang-op-developer |
| synced-from | https://gitcode.com/cann/cannbot-skills |
| synced-date | 2026-05-26 |
| synced-commit | ac5bbd2b4cf427d011874e11f8d1e8b1bef66eda |
| license | UNKNOWN |
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 用法和模式选择
2.2 同类算子示例
生成代码前，必须查阅 examples/ 中的同类算子：
| 算子类型 | 参考示例 |
|---|
| 逐元素运算（add/mul/sigmoid/relu） | examples/elementwise/、examples/activation/ |
| 归约运算（reduce_sum/max/min） | examples/reduce/ |
| 归一化（softmax/layernorm/rmsnorm） | examples/softmax/、examples/normalization/ |
| GEMM | examples/gemm/、examples/developer_mode/gemm_developer.py |
| 融合算子 | examples/flash_attention/ |
| Developer 模式 | examples/developer_mode/ |
查阅示例时关注：
- Kernel 结构：
T.Kernel 参数、cid/vid 用法
- Buffer 分配方式：shape 和 dtype
- pass_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, dtype="float"):
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_npu=True) as (cid, vid):
pass
return main
if __name__ == "__main__":
torch.manual_seed(0)
test_configs = [...]
for config in test_configs:
pass
print("All tests passed!")
步骤 4：运行验证
python examples/{op}/example_{op}.py
如果报错，按以下顺序排查：
- 编译错误 → 检查 buffer 大小、API 参数、对齐
- 运行错误 → 检查索引越界、同步缺失
- 精度错误 → 检查计算公式、数据类型、容差设置
4. 关键编码规范
Buffer 分配
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])
同步
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[...])
pass_configs = {
tilelang.PassConfigKey.TL_ASCEND_AUTO_SYNC: True,
tilelang.PassConfigKey.TL_ASCEND_MEMORY_PLANNING: True,
}
广播
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)
测试模板
ref_output = torch.nn.functional.softmax(input_data, dim=-1)
torch.testing.assert_close(output.cpu(), ref_output.cpu(), rtol=rtol, atol=atol)
5. Checklist
生成代码后逐项检查：
| # | 检查项 |
|---|
| 1 | out_idx 与函数签名中的输出参数位置一致 |
| 2 | block_M // VEC_NUM 在 buffer 分配和索引中一致使用 |
| 3 | 所有 T.alloc_ub 的 shape 乘积不超 UB 容量 |
| 4 | Expert 模式有 T.Scope("V") 和 T.barrier_all() |
| 5 | Developer 模式有对应的 pass_configs |
| 6 | 测试包含至少 2 个配置（小规模 + 典型规模） |
| 7 | golden 函数使用 PyTorch 标准实现 |
