向TileGym添加cuTile内核Skill tilegym-adding-cutile-kernel

名称	tilegym-adding-cutile-kernel
描述	向 TileGym 添加一个新的 cuTile GPU 内核运算符。涵盖 ops.py 中的分发注册、cuTile 后端实现、init.py 导出、测试创建以及 tests/benchmark 中的基准测试。当向 TileGym 添加、创建或实现新的 cuTile 运算符/内核，或询问如何注册新的 cuTile op 时使用。
开源协议	CC-BY-4.0 AND Apache-2.0 metadata:
作者	“TileGym Team TileGym@nvidia.com” tags: - cutile - kernel - tilegym - gpu - dispatch

向 TileGym 添加 cuTile 内核

用于添加一个带有 cuTile 后端的新运算符（例如 my_op）的端到端工作流。

执行规则

必须严格遵守以下规则：

1. 在编写任何代码之前，使用 TodoWrite 创建下面的清单
2. 按顺序执行步骤 — 不要跳过或合并步骤
3. 完成后将每个待办标记为 completed，开始时标记为 in_progress
4. 如果某个步骤不适用（例如，没有 cuTile 实现），用注释将其标记为 completed，不要静默跳过
5. 每个步骤必须产生文件写入或明确的跳过决定 —— 不得有遗漏

说明

必须在开始时将此清单复制到 TodoWrite：

- [ ] 步骤 1：在 ops.py 中注册分发接口
- [ ] 步骤 2：实现 cuTile 后端
- [ ] 步骤 3：在 __init__.py (cutile) 中注册
- [ ] 步骤 4：添加测试
- [ ] 步骤 5：在 tests/benchmark 中添加基准测试
- [ ] 步骤 6：验证（运行 pytest + lint）

步骤 1：注册分发接口

文件：src/tilegym/ops/ops.py

添加一个 @dispatch 函数——这是所有后端的单一入口点。

@dispatch(
    "my_op",
)
def my_op(
    input: torch.Tensor,
    out: Optional[torch.Tensor] = None,
    **kwargs: Any,
):
    """
    Description of my_op.

    Args:
        input: Input tensor
        out: Optional preallocated output tensor
        **kwargs: Additional arguments for backend-specific configurations

    Returns:
        torch.Tensor
    """
    raise NotImplementedError(f"my_op is not implemented for {get_current_backend()}")

关键规则：

• 函数主体仅抛出 NotImplementedError
• 包含 **kwargs 以支持后端特定参数

参考：参见 src/tilegym/ops/ops.py 中的现有操作（例如 silu_and_mul、softmax）

步骤 2：实现 cuTile 后端

文件：src/tilegym/ops/cutile/my_op.py

文件结构遵循以下模板：

import torch
import cuda.tile as ct

from tilegym.backend import register_impl


@ct.kernel
def my_op_kernel_ct(x, output, n_elements: ct.Constant[int], BLOCK_SIZE: ct.Constant[int]):
    bid = ct.bid(0)
    indices = bid * BLOCK_SIZE + ct.arange(0, BLOCK_SIZE)
    x_val = ct.gather(x, indices)
    # ... compute ...
    ct.scatter(output, indices, result)


@register_impl("my_op", backend="cutile")
def my_op(input: torch.Tensor, out: torch.Tensor = None, **kwargs) -> torch.Tensor:
    n = input.numel()
    if out is None:
        out = torch.empty_like(input)
    grid = ((n + 1023) // 1024,)
    ct.launch(stream, grid, kernel, (some args, ...))
    return out

参考：src/tilegym/ops/cutile/silu_and_mul.py

步骤 3：在 __init__.py 中注册（关键）

缺少此步骤意味着 cuTile 后端实现永远不会被加载。

文件：src/tilegym/ops/cutile/__init__.py

在 if is_backend_available("cutile"): 块内按字母顺序添加：

from . import my_op

并且在函数导入部分：

from .my_op import my_op

并将 "my_op" 添加到 __all__。

步骤 4：添加测试

文件：tests/ops/test_my_op.py

关键：始终从 tilegym.ops 导入，永远不要从 tilegym.ops.cutile.my_op 导入。

import pytest
import torch

from tilegym.backend import is_backend_available, set_backend
from .. import common

_backends = ["cutile"]


class Test_MY_OP(common.PyTestCase):
    @staticmethod
    def reference(input):
        """Reference implementation using PyTorch."""
        return torch.some_reference(input)

    @pytest.mark.parametrize("shape, dtype", [
        ((1024,), torch.float16),
        ((1024, 512), torch.float32),
        ((64, 64, 64), torch.bfloat16),
    ])
    @pytest.mark.parametrize("backend", _backends)
    def test_op(self, shape, dtype, backend, arch):
        if backend == "cutile" and not is_backend_available("cutile"):
            pytest.skip("Cutile backend not available")
        try:
            set_backend(backend)
        except Exception as e:
            pytest.skip(f"Backend is not supported: {e}")

        self.setUp()

        from tilegym.ops import my_op

        A = torch.randn(*shape, dtype=dtype, device="cuda")
        self.assertCorrectness(
            my_op, self.reference, {"input": A},
            atol=1e-3, rtol=1e-3,
        )

关键模式：

• _backends = ["cutile"]
• test_op：使用带 try-except 的 set_backend(backend)，调用 self.setUp()

参考：tests/ops/test_silu_and_mul.py

以下是常见错误。

1. 缺少 _backends 列表（在类内部）
2. test_op / test_op_xxx — 缺少 @pytest.mark.parametrize(“backend”, _backends)、backend 参数，以及 tilegym.is_backend_available / tilegym.set_backend 模式

步骤 5：将基准测试添加到 tests/benchmark

文件：tests/benchmark/bench_my_op.py

来自 benchmark_rules.md 的关键规则：

• 通过 tilegym.ops.my_op(a, b, ..., backend=backend) 调用操作——不要使用 set_backend。
• 定义 ALL_BACKENDS（至少包括 cutile 和 torch），使用 get_supported_backends() 进行过滤。
• 实现 reference_my_op(...) 并注册：register_impl("my_op", "torch")(reference_my_op)。
• 使用 create_benchmark_config() 构建 triton.testing.Benchmark 配置（例如按 shape/dtype）。
• 在 bench_my_op(...) 上使用 @triton.testing.perf_report([...])；在基准函数内部：使用 torch.testing.assert_close(fn(), ref(), ...) 进行正确性检查，然后 ms = triton.testing.do_bench(fn)（或 do_bench_cudagraph），计算 GB/s 或 TFLOPS，并返回指标。
• 入口点：if __name__ == "__main__": bench_my_op.run(print_data=True)。

模板结构：

import torch
import triton
import triton.testing

import tilegym
from tilegym.backend import is_backend_available, register_impl

ALL_BACKENDS = [
    ("cutile", "cuTile", ("orange", "-")) if is_backend_available("cutile") else None,
    ("torch", "PyTorch", ("green", "-")),
]

def get_supported_backends():
    return [p for p in ALL_BACKENDS if p is not None]

def reference_my_op(input: torch.Tensor, out: torch.Tensor = None, **kwargs):
    """Reference implementation using PyTorch."""
    ...

register_impl("my_op", "torch")(reference_my_op)

def create_benchmark_config(datatype, ...):
    available_backends = get_supported_backends()
    if not available_backends:
        return None
    backends, names, styles = zip(*available_backends)
    return triton.testing.Benchmark(
        x_names=["M"],  # or other dimension names
        x_vals=[...],
        line_arg="backend",
        line_vals=list(backends),
        line_names=list(names),
        styles=list(styles),
        ylabel="GB/s",  # or TFLOPS
        plot_name="my-op-...",
        args={"datatype": datatype, ...},
    )

@triton.testing.perf_report([
    create_benchmark_config(datatype, ...)
    for datatype in [torch.float16, torch.float32]
    for ... in [...]
])
def bench_my_op(M, backend, datatype, ..., device="cuda"):
    x = torch.randn(..., dtype=datatype, device=device)

    fn = lambda: tilegym.ops.my_op(x, backend=backend)
    ref = lambda: reference_my_op(x)
    torch.testing.assert_close(fn(), ref(), rtol=1e-2, atol=1e-2)

    ms = triton.testing.do_bench(fn)  # or do_bench_cudagraph(fn)
    # Compute metric (e.g. GB/s or TFLOPS) from ms and problem size
    return metric

if __name__ == "__main__":
    bench_my_op.run(print_data=True)

基准图名称：必须包含 -TFLOPS 或 -GBps 后缀

• 示例：plot_name=f"persistent-layer-norm-M{num_rows}-{dtype_name}-GBps"

步骤 6：验证

# Run tests
pytest tests/ops/test_my_op.py -v

# Run benchmark (optional)
python tests/benchmark/bench_my_op.py

# Lint
pre-commit run -a