Cómo añadir un kernel de cuTile a TileGym

Flujo de trabajo completo para añadir un nuevo operador (por ejemplo, my_op) con el backend cuTile.

Normas de ejecución

DEBES seguir estas reglas al pie de la letra:

1. Utiliza TodoWrite para crear la lista de comprobación que aparece a continuación ANTES de escribir ningún código
2. Ejecuta los pasos en orden: NO te saltes ninguno ni combines pasos
3. Marca cada tarea como «completada» al terminarla, y como «en curso » al comenzarla
4. Si un paso no es aplicable (por ejemplo, si no hay una implementación de cuTile), márcalo como «completado» con una nota; NO lo omitas sin avisar
5. Cada paso DEBE dar lugar a la escritura de un archivo o a una decisión explícita de omitirlo; no se admiten omisiones silenciosas

Instrucciones

DEBES copiar esta lista de comprobación en TodoWrite al principio:

- [ ] Paso 1: Registrar la interfaz de despacho en ops.py
- [ ] Paso 2: Implementar el backend de cuTile
- [ ] Paso 3: Registrarlo en __init__.py (cutile)
- [ ] Paso 4: Añadir pruebas
- [ ] Paso 5: Añadir el benchmark a tests/benchmark
- [ ] Paso 6: Verificar (ejecutar pytest + lint)

Paso 1: Registrar la interfaz de despacho

Archivo: src/tilegym/ops/ops.py

Añadir una función @dispatch: este es el único punto de entrada para todos los backends.

@dispatch(
    "my_op",
)
def my_op(
    input: torch.Tensor,
    out: Optional[torch.Tensor] = None,
    **kwargs: Any,
):
    """
    Descripción de my_op.

    Argumentos:
        input: Tensor de entrada
        out: Tensor de salida preasignado opcional
        **kwargs: Argumentos adicionales para configuraciones específicas del backend

    Devuelve:
        torch.Tensor
    """
    raise NotImplementedError(f"my_op no está implementado para {get_current_backend()}")

Reglas clave:

• El cuerpo de la función solo lanza un NotImplementedError
• Incluye **kwargs para los parámetros específicos del backend

Referencia: Consulta las operaciones existentes en src/tilegym/ops/ops.py (p. ej., silu_and_mul, softmax)

Paso 2: Implementar el backend cuTile

Archivo: src/tilegym/ops/cutile/my_op.py

La estructura del archivo sigue esta plantilla:

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)
    # ... cálculo ...
    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

Referencia: src/tilegym/ops/cutile/silu_and_mul.py

Paso 3: Registrarse en __init__.py (FUNDAMENTAL)

Si se omite este paso, la implementación del backend cuTile nunca se cargará.

Archivo: src/tilegym/ops/cutile/__init__.py

Añadir dentro del bloque if is_backend_available("cutile"): (por orden alfabético):

from . import my_op

Y en la sección de importación de funciones:

from .my_op import my_op

Y añade «my_op» a __all__.

Paso 4: Añadir pruebas

Archivo: tests/ops/test_my_op.py

IMPORTANTE: Importa siempre desde tilegym.ops, NUNCA desde 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):
        """Implementación de referencia utilizando 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("El backend Cutile no está disponible")
        try:
            set_backend(backend)
        except Exception as e:
            pytest.skip(f"El backend no es compatible: {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,
        )

Patrones clave:

• _backends = ["cutile"]
• test_op: utiliza set_backend(backend) con try-except, llama a self.setUp()

Referencia: tests/ops/test_silu_and_mul.py

A continuación se indican los errores más comunes.

1. Falta la lista _backends (dentro de la clase)
2. test_op / test_op_xxx: faltan @pytest.mark.parametrize("backend", _backends), el parámetro backend y el patrón tilegym.is_backend_available / tilegym.set_backend

Paso 5: Añadir la prueba de rendimiento a tests/benchmark

Archivo: tests/benchmark/bench_my_op.py

Reglas clave de benchmark_rules.md:

• Llama a la operación mediante tilegym.ops.my_op(a, b, ..., backend=backend); no utilices set_backend.
• Define ALL_BACKENDS (incluye al menos cutile y torch), y filtra con get_supported_backends().
• Implementa reference_my_op(...) y regístrala: register_impl("my_op", "torch")(reference_my_op).
• Utiliza create_benchmark_config() para crear configuraciones de triton.testing.Benchmark (por ejemplo, por shape/dtype).
• Utiliza @triton.testing.perf_report([...]) en bench_my_op(...); dentro de la función de benchmark: comprueba la corrección con torch.testing.assert_close(fn(), ref(), ...), luego ms = triton.testing.do_bench(fn) (o do_bench_cudagraph), calcula GB/s o TFLOPS y devuelve la métrica.
• Punto de entrada: if __name__ == "__main__": bench_my_op.run(print_data=True).

Estructura de la plantilla:

import torch
import triton
import triton.testing

import tilegym
from tilegym.backend import is_backend_available, register_impl

ALL_BACKENDS = [
    ("cutile", "cuTile", ("orange", "-")) si is_backend_available("cutile"); en caso contrario, 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):
    """Implementación de referencia utilizando 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"],  # u otros nombres de dimensiones
        x_vals=[...],
        line_arg="backend",
        line_vals=list(backends),
        line_names=list(names),
        styles=list(styles),
        ylabel="GB/s",  # o 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)  # o do_bench_cudagraph(fn)
    # Calcular la métrica (p. ej., GB/s o TFLOPS) a partir de ms y del tamaño del problema
    return metric

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

Nombres de los gráficos de referencia: deben incluir el sufijo -TFLOPS o -GBps

• Ejemplo: plot_name=f"persistent-layer-norm-M{num_rows}-{dtype_name}-GBps"

Paso 6: Verificar

# Ejecutar pruebas
pytest tests/ops/test_my_op.py -v

# Ejecutar el benchmark (opcional)
python tests/benchmark/bench_my_op.py

# Lint
pre-commit run -a