Tilegym Adding Cutile Kernel

@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()}")

import torchimport cuda.tile as ct
from tilegym.backend import register_impl
@ct.kerneldef 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

from . import my_op

from .my_op import my_op

import pytestimport torch
from tilegym.backend import is_backend_available, set_backendfrom .. 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,        )

import torchimport tritonimport triton.testing
import tilegymfrom 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)