Tilelang Example Merge

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" 等命名模式

输入

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

• 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 ...）。

例如 test_softmax.py 的 golden 是：

def golden_softmax(x):
    return torch.nn.functional.softmax(x.float(), dim=-1).to(x.dtype)

复制到 example_softmax.py 后保留为独立函数，循环内 ref = golden_softmax(x)。

精度阈值提取

精度阈值必须根据 test_configs 中选中用例的 dtype 动态确定，不能硬编码某个 dtype 的阈值。提取步骤：

1. 确定选中 dtype：读取选中的 L0/L1 用例的 dtype 字段。
2. 查 test_{op}.py 的 get_precision 表：找到该 dtype 对应的 (atol, rtol, max_abs_limit, required_ratio) 四元组。
3. 填入模板占位符：将四个数值替换模板中的 {atol}、{rtol}、{max_abs_limit}、 {required_ratio}。

各 dtype 的标准阈值参考（源自 test_{op}.py 的 get_precision，与 tilelang-op-test-design/references/precision-standard.md 一致）：

注意：不同算子的 test_{op}.py 可能只覆盖表中部分 dtype。以目标算子 test 文件 中实际存在的为准，不要套用上表缺失的 dtype。

多 dtype 处理：如果选中用例存在多个不同 dtype（少见），不能在循环外写死一组 阈值。需在循环内按 dtype 分支查表，例如：

for B, N, block_M, block_N, dtype, level in test_configs:
    ...
    # Precision thresholds by dtype (from test_{op}.py get_precision)
    if dtype == "float16":
        atol, rtol, max_abs_limit, required_ratio = 2**-14, 2**-9, 1e-1, 0.99
    elif dtype == "float32" or dtype == "float":
        atol, rtol, max_abs_limit, required_ratio = 2**-16, 2**-10, 1e-2, 0.99
    # ... 只列选中用例实际涉及的 dtype
    ratio = (abs_err <= (atol + rtol * ref_cpu[m].abs())).float().mean().item()
    max_abs = abs_err.max().item()
    assert ratio >= required_ratio and max_abs <= max_abs_limit, ...

单 dtype（常见）：如果选中用例 dtype 相同，直接在循环外写死该 dtype 的四个数值 字面量（不分支、不封装函数），保持代码精简。

第四步：生成合并文件

按以下模板生成 example_{op}.py（参考 examples/normalization/layer_norm.py 和 examples/developer_mode/gelu_mul_developer.py 的仓库惯例）：

import tilelang
from tilelang import language as T
import torch

tilelang.cache.clear_cache()

# ========== Operator Implementation ==========
# （从 {op}.py 复制的 pass_configs、常量、@tilelang.jit 函数，原样保留）

pass_configs = {
    tilelang.PassConfigKey.TL_ASCEND_AUTO_SYNC: True,
    # ... 其他配置
}

@tilelang.jit(out_idx=[1], pass_configs=pass_configs)
def {op}(...):
    """{算子简述}"""
    # ... kernel 完整实现 ...
    return main


# ========== Golden reference ==========
def golden_{op}(x):
    """{一句话说明}"""
    return {test_{op}.py 中的 golden 函数体，原样复制}


# ========== Tests ==========
if __name__ == "__main__":
torch.manual_seed(0)

# Representative configs: 1 L0 + 1 L1
test_configs = [
    # (B, N, block_M, block_N, dtype, level)
    (..., ..., ..., ..., "...", "L0"),  # {L0 选中用例描述}
    (..., ..., ..., ..., "...", "L1"),  # {L1 选中用例描述}
]

for B, N, block_M, block_N, dtype, level in test_configs:
    print(f"Testing {op} {level} with B={B}, N={N}, block=({block_M},{block_N}), dtype={dtype}")
    func = {op}(B, N, block_M, block_N, dtype=dtype)
    print("Init successful!")
    torch_dtype = getattr(torch, dtype) if dtype != "float" else torch.float32
    x = torch.randn(B, N, dtype=torch_dtype).npu()
    y = func(x)
    ref = golden_{op}(x)
    # Precision check ({dtype} mixed tolerance, inlined — thresholds from test_{op}.py)
    y_cpu, ref_cpu = y.detach().cpu().float(), ref.detach().cpu().float()
    m = torch.isfinite(ref_cpu)
    abs_err = (y_cpu[m] - ref_cpu[m]).abs()
    ratio = (abs_err <= ({atol} + {rtol} * ref_cpu[m].abs())).float().mean().item()
    max_abs = abs_err.max().item()
    assert ratio >= {required_ratio} and max_abs <= {max_abs_limit}, f"precision fail: ratio={ratio:.4f} max_abs={max_abs:.3e}"
    print(f"Test pass! matched_ratio={ratio:.4f} max_abs={max_abs:.3e}")

print("Kernel Output Match!")

关键格式约定（必须遵循仓库惯例）

1. 精度检查：内联混合容差检查，不要用 torch.testing.assert_close，也不要 封装成函数。直接在测试用例中按选中 dtype 的阈值内联计算。阈值根据 test_configs 中选中用例的 dtype 从 test_{op}.py 的 get_precision 表动态查取，填入模板的 {atol}/{rtol}/{max_abs_limit}/{required_ratio} 占位符——禁止硬编码某个 固定 dtype 的阈值。双门控：逐元素 |actual-golden| <= atol + rtol*|golden|， 整体 matched_ratio >= required_ratio 且 max_abs_error <= max_abs_limit。查表 规则和多 dtype 处理见上方"精度阈值提取"小节。
2. 缓存清理：文件顶部用 tilelang.cache.clear_cache()（不是 tilelang.disable_cache()）， 与仓库现有示例一致。
3. 输入数据：用 torch.randn(...).npu() 生成随机输入。如果原测试用了特定数值范围 （如 uniform_(-1000, 1000)），L1 代表用例可以保留该范围，但 L0 用标准 randn。
4. 打印格式：print(f"Testing {op} ... with ...") → print("Init successful!") → print("Test pass!") → 末尾 print("Kernel Output Match!")，与仓库现有示例一致。
5. 无 import 兄弟模块：example_{op}.py 中禁止出现 from {op} import {op} 或 sys.path.insert 等导入语句。kernel 代码直接内联。
6. 无分层测试框架：不要保留 --level 参数分发、COVERAGE_CATEGORY、 L1_CASES 列表、check_precision 等分层测试基础设施。上库示例用单个 test_configs 列表 + for 循环顺序执行两个代表性用例（1 个 L0 + 1 个 L1）， 循环体内复用同一套 kernel 编译/运行/golden/精度检查逻辑，避免代码重复。 这与仓库现有示例（layer_norm.py、gelu_mul_developer.py 的 test_configs 循环）一致。
7. 使用 if __name__ == "__main__" 守卫：测试代码放在 if __name__ == "__main__": 块中，python example_{op}.py 直接运行时会执行两个代表性用例并打印 PASS/FAIL。 这比仓库现有示例（layer_norm.py 等用模块级测试）更显式，且 import 时不自动执行。

第五步：验证

生成文件后，运行验证：

source set_env.sh
python examples/{op}/example_{op}.py

确认输出包含 "Kernel Output Match!"。如果失败，检查：

• kernel 代码是否完整复制（漏了常量或辅助函数）
• golden 实现是否正确
• shape/dtype 是否与原测试一致
• 是否有遗留的 import 语句

输出文件结构

生成的 example_{op}.py 分四段：

1. imports + tilelang.cache.clear_cache()
2. kernel 实现（pass_configs + @tilelang.jit 函数）   ← 从 {op}.py 复制
3. golden 参考实现（独立函数）                          ← 从 test_{op}.py 复制
4. if __name__ == "__main__": 测试代码（test_configs 循环 + 末尾打印）  ← 精简

目标行数：通常 80-150 行（kernel 行数 + 每个测试约 10-15 行）。

注意事项

• 不要修改 kernel 逻辑：kernel 代码从 {op}.py 原样复制，不做任何改动。如果 kernel 依赖模块级辅助函数（如 cast_or_copy），一并复制。
• 保留必要的模块级常量：pass_configs、CAST_MODE_*、VEC_NUM 等被 kernel 使用的常量必须保留。未被选中队列测试使用的常量（如 COVERAGE_MANIFEST）丢弃。
• golden 简化但不失真：golden 函数保留正确的数学逻辑，但可以去掉冗长的 docstring。 如果原 golden 调用了 PyTorch 内置函数（如 F.softmax、torch.layer_norm），直接用 该调用，不要重新实现。
• dtype 处理：如果 kernel 的 dtype 参数用 "float" 表示 float32，测试中需用 getattr(torch, dtype) if dtype != "float" else torch.float32 转换，与仓库惯例一致。
• 原文件保留：生成 example_{op}.py 后，原 {op}.py 和 test_{op}.py 不删除， 它们仍用于开发阶段的分层测试。example_{op}.py 是上库用的精简单文件。

常见问题

测试文件结构不是标准 L0/L1 分层怎么办？

有些算子的测试文件可能用不同的结构（如单个 test_configs 列表无 L0/L1 区分）。 此时：

• 将第一个规则 shape 用例作为 "L0 representative"
• 将第二个规则 shape 用例（或稍大 shape 的用例）作为 "L1 representative"
• 在注释中标注 "representative" 而非 "L0"/"L1"

kernel 文件有 if __name__ == "__main__" 块怎么办？

纯 kernel 文件（{op}.py）通常没有 __main__ 块。如果有，只复制 kernel 部分 （imports + pass_configs + @tilelang.jit 函数），丢弃 __main__ 块。

算子有多个 kernel 函数怎么办？

如果 {op}.py 包含多个 @tilelang.jit 函数，全部保留（它们可能互相调用或用于 不同配置）。测试代码中调用主 kernel。

生成的文件跑不过怎么办？

最常见原因：

1. 漏复制常量/辅助函数：检查 kernel 是否引用了未复制的模块级符号
2. golden 与 kernel dtype 路径不一致：确保 golden 的 dtype 转换与 kernel 对齐
3. shape 不匹配：确保测试 shape 与 kernel 的 jit 参数一致
4. 缺少 torch.npu.synchronize()：如果原测试有同步调用，保留它

参考示例

仓库中现有的单文件示例（合并后的目标格式参考）：

• examples/normalization/layer_norm.py — kernel + test_configs 循环 + assert_close
• examples/developer_mode/gelu_mul_developer.py — kernel + test_configs 循环 + assert_close
• examples/normalization/rms_norm.py — 同上模式

这些文件的共同特征：单文件、模块级测试、tilelang.cache.clear_cache()、末尾 print("Kernel Output Match!")。

与本 skill 的差异：仓库现有示例用 torch.testing.assert_close(rtol=1e-2, atol=1e-2) 做精度检查，但本 skill 按用户要求采用 test_{op}.py 的混合容差标准（内联，按 dtype 写死阈值），比 assert_close 的单一 rtol/atol 更贴合算子精度分级要求。其余格式 （单文件、clear_cache、末尾打印）保持与仓库惯例一致。