Hinzufügen eines cuTile-Kernels zu TileGym

End-to-End-Workflow zum Hinzufügen eines neuen Operators (z. B. my_op) mit cuTile-Backend.

Ausführungsregeln

Diese Regeln MÜSSEN strikt befolgt werden:

1. Verwenden Sie „TodoWrite“, um die unten stehende Checkliste zu erstellen, BEVOR Sie Code schreiben
2. Führen Sie die Schritte der Reihe nach aus – überspringen Sie KEINE Schritte und kombinieren Sie sie NICHT
3. Markieren Sie jede Aufgabe nach Abschluss als „abgeschlossen“ und zu Beginn als „in Bearbeitung“
4. Wenn ein Schritt nicht zutrifft (z. B. keine cuTile-Implementierung), markieren Sie ihn mit einem Hinweis als „abgeschlossen“ und überspringen Sie ihn NICHT stillschweigend
5. Jeder Schritt MUSS zu einem Dateischreibvorgang oder einer expliziten Entscheidung zum Überspringen führen – keine stillschweigenden Auslassungen

Anweisungen

Diese Checkliste MUSS zu Beginn in „TodoWrite“ kopiert werden:

- [ ] Schritt 1: Dispatch-Schnittstelle in ops.py registrieren
- [ ] Schritt 2: cuTile-Backend implementieren
- [ ] Schritt 3: In __init__.py (cutile) registrieren
- [ ] Schritt 4: Tests hinzufügen
- [ ] Schritt 5: Benchmark zu tests/benchmark hinzufügen
- [ ] Schritt 6: Überprüfen (pytest + lint ausführen)

Schritt 1: Dispatch-Schnittstelle registrieren

Datei: src/tilegym/ops/ops.py

Fügen Sie eine @dispatch -Funktion hinzu – dies ist der einzige Einstiegspunkt für alle Backends.

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

    Argumente:
        input: Eingabetensor
        out: Optionaler, vorab zugewiesener Ausgabetensor
        **kwargs: Zusätzliche Argumente für backend-spezifische Konfigurationen

    Rückgabewert:
        torch.Tensor
    """
    raise NotImplementedError(f"my_op ist für {get_current_backend()} nicht implementiert")

Wichtige Regeln:

• Der Funktionskörper löst nur einen NotImplementedError aus
• Fügen Sie **kwargs für backend-spezifische Parameter ein

Referenz: Siehe vorhandene Operationen in src/tilegym/ops/ops.py (z. B. silu_and_mul, softmax)

Schritt 2: cuTile-Backend implementieren

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

Die Dateistruktur folgt dieser Vorlage:

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

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

Schritt 3: Registrierung in __init__.py (WICHTIG)

Wird dieser Schritt übersprungen, wird die cuTile-Backend-Implementierung niemals geladen.

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

Fügen Sie innerhalb des Blocks `if is_backend_available("cutile")` Folgendes ein ( in alphabetischer Reihenfolge):

from . import my_op

Und im Abschnitt „import“ der Funktion:

from .my_op import my_op

Und fügen Sie „my_op“ zu __all__ hinzu.

Schritt 4: Tests hinzufügen

Datei: tests/ops/test_my_op.py

WICHTIG: Importiere immer aus `tilegym.ops`, NIEMALS aus `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):
        """Referenzimplementierung mit 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 nicht verfügbar")
        try:
            set_backend(backend)
        except Exception as e:
            pytest.skip(f"Backend wird nicht unterstützt: {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,
        )

Schlüsselmuster:

• _backends = ["cutile"]
• test_op: Verwende set_backend(backend) mit try-except, rufe self.setUp() auf

Referenz: tests/ops/test_silu_and_mul.py

Nachfolgend sind die häufigsten Fehler aufgeführt.

1. Fehlende _backends-Liste (innerhalb der Klasse)
2. test_op / test_op_xxx – fehlendes @pytest.mark.parametrize("backend", _backends), fehlender „backend“-Parameter sowie fehlendes Muster tilegym.is_backend_available / tilegym.set_backend

Schritt 5: Füge einen Benchmark zu „tests/benchmark“ hinzu

Datei: tests/benchmark/bench_my_op.py

Wichtige Regeln aus „benchmark_rules.md“:

• Rufe die Operation über `tilegym.ops.my_op(a, b, ..., backend=backend) ` auf – verwende nicht `set_backend`.
• Definieren Sie ALL_BACKENDS (einschließlich mindestens cutile und torch) und filtern Sie mit get_supported_backends().
• Implementieren Sie `reference_my_op(...)` und registrieren Sie es: `register_impl("my_op", "torch")(reference_my_op)`.
• Verwende create_benchmark_config(), um triton.testing.Benchmark -Konfigurationen zu erstellen (z. B. nach shape/dtype).
• Verwenden Sie @triton.testing.perf_report([...]) auf bench_my_op(...); innerhalb der „bench“-Funktion: Korrektheitsprüfung mit `torch.testing.assert_close(fn(), ref(), ...)`, anschließend `ms = triton.testing.do_bench(fn) ` (oder `do_bench_cudagraph`), Berechnung von GB/s oder TFLOPS und Rückgabe der Metrik.
• Einstiegspunkt: if __name__ == "__main__": bench_my_op.run(print_data=True).

Vorlagenstruktur:

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):
    """Referenzimplementierung mit 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"],  # oder andere Dimensionsnamen
        x_vals=[...],
        line_arg="backend",
        line_vals=list(backends),
        line_names=list(names),
        styles=list(styles),
        ylabel="GB/s",  # oder 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)  # oder do_bench_cudagraph(fn)
    # Berechne die Metrik (z. B. GB/s oder TFLOPS) aus ms und der Problemgröße
    return metric

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

Namen für Benchmark-Diagramme: Müssen die Endung -TFLOPS oder -GBps enthalten

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

Schritt 6: Überprüfen

# Tests ausführen
pytest tests/ops/test_my_op.py -v

# Benchmark ausführen (optional)
python tests/benchmark/bench_my_op.py

# Lint
pre-commit run -a