tilelang-ascend: Skill for Claude Code

TileLang Example Merge

概述

将算子开发阶段的双文件结构（{op}.py 纯 kernel + test_{op}.py 分层测试套件） 合并为仓库上库用的单文件 example_{op}.py。

为什么要合并：算子开发时 kernel 和测试分离便于迭代，但仓库上库只接收单文件 示例（参考 examples/normalization/layer_norm.py、examples/developer_mode/gelu_mul_developer.py 的惯例）。单文件示例自包含、可直接 python example_{op}.py 运行验证。

合并策略：kernel 完整保留 + 从测试套件中自动选取 1 个 L0 代表性用例 + 1 个 L1 代表性用例，精简辅助函数，使用 torch.testing.assert_close 做精度检查。

触发条件

• 用户提到"上库"、"提交 PR"、"合并文件"、"生成 example"、"单文件提交"
• 用户要把算子代码整理成仓库可接收的单文件示例
• 用户提到 "example_softmax.py"、"example_layer_norm.py" 等命名模式

输入

参数	说明	示例
算子名	算子目录名和文件名前缀	softmax

输入文件（隐式从算子名推导）：

• examples/{op}/{op}.py — 纯 kernel 文件
• examples/{op}/test_{op}.py — 分层测试文件

输出文件：

• examples/{op}/example_{op}.py — 合并后的单文件示例

工作流程

第一步：读取源文件

1. 确认算子名（用户指定或从对话上下文提取）
2. 读取 examples/{op}/{op}.py，提取完整 kernel 代码
   • 包括模块级常量（pass_configs、CAST_MODE_* 等）
   • 包括 @tilelang.jit 装饰的函数及其内部的 @T.prim_func
   • 不要包含 if __name__ == "__main__" 块（如果有的话）
3. 读取 examples/{op}/test_{op}.py，理解测试结构
   • 识别 L0 测试用例（通常在 test_{op}_l0() 函数或 test_configs 列表中）
   • 识别 L1 测试用例（通常在 test_{op}_l1() 函数或 L1_CASES 列表中）
   • 提取 golden 参考实现函数
   • 提取 get_precision 函数及 dtype→阈值映射表（用于第四步查表填占位符）

第二步：选取代表性用例

L0 代表性用例选取

L0 是门槛测试（规则 shape，block 整除），选取最具代表性的一个：

1. 优先：名称含 "typical" 或 "standard" 的用例（如 l0_typical）
2. 次选：shape 最大的用例（最大 N 或最大 B×N，最能代表真实工作负载）
3. 兜底：第一个 L0 用例

L1 代表性用例选取

L1 是功能测试（含不规则 shape、数值范围覆盖），选取最标准的规则 shape 用例：

1. 优先：带 D-SHAPE-ALIGNED tag 的用例（规则 shape，无尾块）
2. 次选：第一个 shape 为规则对齐的用例（B % block_M == 0 且 N % block_N == 0）
3. 兜底：第一个 L1 用例

选取时注意避开极端边界用例（如 B=1、N=1、超大数值范围），这些适合分层测试 但不适合作为上库示例的代表用例。

第三步：提取 golden 参考实现

从 test_{op}.py 中提取 golden 函数（通常名为 golden_{op} 或直接内联在测试中）， 作为独立函数复制到 example_{op}.py 的 kernel 之后、if __name__ 块之前。

提取规则：

1. 保留独立函数：不要内联到测试循环里。golden 函数放循环外，循环内调用 ref = golden(x)。这样 golden 逻辑只写一遍，多个用例复用，与 test_{op}.py 结构一致。
2. 原样复制函数体：保留数学逻辑，去掉冗长 docstring（一行注释说明即可）。
3. 不要重新实现：如果原 golden 调用了 PyTorch 内置函数（如 F.softmax、 torch.layer_norm），直接用该调用，不要手写等价实现，避免引入新 bug。
4. 函数签名对齐：golden 函数的输入参数应与测试循环中传入的张量一致（通常是 def golden(x): return ...）。