Ajout d'un noyau cuTile à TileGym

Workflow de bout en bout pour l'ajout d'un nouvel opérateur (par exemple, my_op) avec le backend cuTile.

Règles d'exécution

Vous DEVEZ respecter strictement ces règles :

1. Utilisez TodoWrite pour créer la liste de contrôle ci-dessous AVANT d'écrire le moindre code
2. Exécutez les étapes dans l’ordre — NE SAUTEZ PAS d’étapes et NE LES COMBINEZ PAS
3. Marquez chaque tâche comme « terminée » une fois terminée, ou « en cours » lorsque vous la commencez
4. Si une étape n’est pas applicable (par exemple, pas d’implémentation de cuTile), marquez-la comme « terminée » avec une note, ne la sautez PAS sans le signaler
5. Chaque étape DOIT aboutir à l'écriture d'un fichier ou à une décision explicite de sauter l'étape — aucune omission silencieuse

Instructions

Vous DEVEZ copier cette liste de contrôle dans TodoWrite au début :

- [ ] Étape 1 : Enregistrer l’interface de répartition dans ops.py
- [ ] Étape 2 : Implémenter le backend cuTile
- [ ] Étape 3 : Enregistrer dans __init__.py (cutile)
- [ ] Étape 4 : Ajouter des tests
- [ ] Étape 5 : Ajouter un benchmark dans tests/benchmark
- [ ] Étape 6 : Vérifier (lancer pytest + lint)

Étape 1 : Enregistrer l'interface de répartition

Fichier: src/tilegym/ops/ops.py

Ajouter une fonction @dispatch — il s'agit du point d'entrée unique pour tous les backends.

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

    Arguments :
        input : Tenseur d’entrée
        out : Tenseur de sortie préalloué (Optional)
        **kwargs : Arguments supplémentaires pour les configurations spécifiques au backend

    Renvoie :
        torch.Tensor
    """
    raise NotImplementedError(f"my_op n'est pas implémenté pour {get_current_backend()}")

Règles clés :

• Le corps de la fonction ne lève qu’ une exception NotImplementedError
• Inclure **kwargs pour les paramètres spécifiques au backend

Référence: voir les opérations existantes dans src/tilegym/ops/ops.py (par exemple, silu_and_mul, softmax)

Étape 2 : Implémenter le backend cuTile

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

La structure du fichier suit ce modèle :

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)
    # ... calcul ...
    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

Référence: src/tilegym/ops/cutile/silu_and_mul.py

Étape 3 : Enregistrement dans __init__.py (CRITIQUE)

Si cette étape est omise, l’implémentation du backend cuTile ne sera jamais chargée.

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

Ajouter à l'intérieur du bloc if is_backend_available("cutile") : (par ordre alphabétique) :

from . import my_op

Et dans la section d'importation des fonctions :

from .my_op import my_op

Et ajoutez « my_op » à __all__.

Étape 4 : Ajouter des tests

Fichier: tests/ops/test_my_op.py

IMPORTANT: Importez toujours depuis tilegym.ops, JAMAIS depuis 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):
        """Implémentation de référence utilisant 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("Backend Cutile indisponible")
        try:
            set_backend(backend)
        except Exception as e:
            pytest.skip(f"Le backend n'est pas pris en charge : {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,
        )

Modèles clés :

• _backends = ["cutile"]
• test_op: utilisez set_backend(backend) avec try-except, appelez self.setUp()

Référence: tests/ops/test_silu_and_mul.py

Voici les erreurs courantes.

1. Liste _backends manquante (à l’intérieur de la classe)
2. test_op / test_op_xxx — absence de @pytest.mark.parametrize("backend", _backends), du paramètre backend et du modèle tilegym.is_backend_available / tilegym.set_backend

Étape 5 : Ajouter un benchmark à tests/benchmark

Fichier: tests/benchmark/bench_my_op.py

Règles clés tirées du fichier benchmark_rules.md :

• Appelez l’opération via tilegym.ops.my_op(a, b, ..., backend=backend) — n’ utilisez pas set_backend.
• Définissez ALL_BACKENDS (incluez au moins cutile et torch), puis filtrez à l’aide de get_supported_backends().
• Implémentez reference_my_op(...) et enregistrez-la : register_impl("my_op", "torch")(reference_my_op).
• Utilisez create_benchmark_config() pour créer des configurations triton.testing.Benchmark (par exemple par shape/dtype).
• Utilisez @triton.testing.perf_report([...]) sur bench_my_op(...); à l’intérieur de la fonction de benchmark : vérifiez l’exactitude avec `torch.testing.assert_close(fn(), ref(), ...)`, puis `ms = triton.testing.do_bench(fn) ` (ou `do_bench_cudagraph`), calculez les GB/s ou TFLOPS, et renvoyez la métrique.
• Point d’entrée : if __name__ == "__main__" : bench_my_op.run(print_data=True).

Structure du modèle :

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):
    """Implémentation de référence utilisant 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"],  # ou d’autres noms de dimensions
        x_vals=[...],
        line_arg="backend",
        line_vals=list(backends),
        line_names=list(names),
        styles=list(styles),
        ylabel="GB/s",  # ou 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)  # ou do_bench_cudagraph(fn)
    # Calculer la métrique (par ex. GB/s ou TFLOPS) à partir de ms et de la taille du problème
    return metric

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

Noms des graphiques de benchmark: doivent inclure le suffixe -TFLOPS ou -GBps

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

Étape 6 : Vérification

# Exécuter les tests
pytest tests/ops/test_my_op.py -v

# Exécuter le benchmark (facultatif)
python tests/benchmark/bench_my_op.py

# Lint
pre-commit run -a