feat: add TileLang copy mode templates

Co-authored-by: wawahejun <hejunlbbc@gmail.com>

examples/05_tilelang_copy_modes.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT / "python"))
+sys.path.insert(0, str(ROOT))
+
+import torch
+
+from ops.common.tilelang.eager_copy import copy_eager, copy_eager_
+from ops.common.tilelang.lazy_out_idx_copy import copy_lazy_out_idx, copy_lazy_out_idx_
+
+
+def main() -> None:
+    src = torch.randn((1024,), device="cuda", dtype=torch.float32)
+
+    eager = copy_eager(src)
+    torch.testing.assert_close(eager, src)
+
+    eager_out = torch.empty_like(src)
+    copy_eager_(eager_out, src)
+    torch.testing.assert_close(eager_out, src)
+
+    lazy = copy_lazy_out_idx(src)
+    torch.testing.assert_close(lazy, src)
+
+    lazy_out = torch.empty_like(src)
+    copy_lazy_out_idx_(lazy_out, src)
+    torch.testing.assert_close(lazy_out, src)
+
+    print("tilelang eager and lazy out_idx copy ok")
+
+
+if __name__ == "__main__":
+    main()

ops/common/tilelang/eager_copy.py

@@ -0,0 +1,86 @@
+from dataclasses import dataclass
+from functools import lru_cache
+
+import tilelang
+import tilelang.language as T
+import torch
+
+
+def _tl_dtype(dtype: torch.dtype):
+    if dtype is torch.float16:
+        return T.float16
+    if dtype is torch.float32:
+        return T.float32
+    raise TypeError(f"unsupported TileLang dtype: {dtype}")
+
+
+@tilelang.jit
+def _copy_eager_kernel(src, BLOCK_N: int, dtype):
+    N = T.const("N")
+    src: T.Tensor((N,), dtype)
+    out = T.empty((N,), dtype)
+
+    with T.Kernel(N // BLOCK_N, threads=256) as pid_n:
+        T.copy(
+            src[pid_n * BLOCK_N : (pid_n + 1) * BLOCK_N],
+            out[pid_n * BLOCK_N : (pid_n + 1) * BLOCK_N],
+        )
+
+    return out
+
+
+def _block_n(n: int) -> int:
+    return 1024 if n % 1024 == 0 else n
+
+
+@lru_cache(maxsize=32)
+def _compiled_copy_eager(n: int, block_n: int, dtype: torch.dtype):
+    return _copy_eager_kernel.compile(N=n, BLOCK_N=block_n, dtype=_tl_dtype(dtype))
+
+
+@dataclass
+class EagerCopyPrepared:
+    out: torch.Tensor
+    src: torch.Tensor
+    kernel: object
+
+    def run(self) -> None:
+        self.out.copy_(self.kernel(self.src))
+
+    def destroy(self) -> None:
+        pass
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        self.destroy()
+
+
+def prepare_copy_eager(out: torch.Tensor, src: torch.Tensor) -> EagerCopyPrepared:
+    if out.shape != src.shape:
+        raise ValueError("eager copy expects matching shapes")
+    if out.dtype != src.dtype:
+        raise TypeError("eager copy expects matching dtypes")
+    if not out.is_cuda or not src.is_cuda:
+        raise ValueError("eager copy expects CUDA tensors")
+    if not out.is_contiguous() or not src.is_contiguous():
+        raise ValueError("eager copy v1 supports contiguous tensors only")
+
+    n = src.numel()
+    block_n = _block_n(n)
+    if n % block_n != 0:
+        raise ValueError("eager copy v1 requires N % BLOCK_N == 0")
+    kernel = _compiled_copy_eager(n, block_n, src.dtype)
+    return EagerCopyPrepared(out, src, kernel)
+
+
+def copy_eager_(out: torch.Tensor, src: torch.Tensor) -> torch.Tensor:
+    with prepare_copy_eager(out, src) as prepared:
+        prepared.run()
+    return out
+
+
+def copy_eager(src: torch.Tensor) -> torch.Tensor:
+    out = torch.empty_like(src)
+    return copy_eager_(out, src)

ops/common/tilelang/lazy_out_idx_copy.py

@@ -0,0 +1,91 @@
+from dataclasses import dataclass
+from functools import lru_cache
+
+import tilelang
+import tilelang.language as T
+import torch
+
+
+def _tl_dtype_str(dtype: torch.dtype) -> str:
+    if dtype is torch.float16:
+        return "float16"
+    if dtype is torch.float32:
+        return "float32"
+    raise TypeError(f"unsupported TileLang dtype: {dtype}")
+
+
+@tilelang.jit(out_idx=[1])
+def _copy_lazy_out_idx_kernel(n: int, block_n: int, dtype: str):
+    @T.prim_func
+    def main(src: T.Tensor((n,), dtype), out: T.Tensor((n,), dtype)):
+        with T.Kernel(n // block_n, threads=256) as pid_n:
+            T.copy(
+                src[pid_n * block_n : (pid_n + 1) * block_n],
+                out[pid_n * block_n : (pid_n + 1) * block_n],
+            )
+
+    return main
+
+
+def _block_n(n: int) -> int:
+    return 1024 if n % 1024 == 0 else n
+
+
+@lru_cache(maxsize=32)
+def _compiled_copy_lazy_out_idx(n: int, block_n: int, dtype: torch.dtype):
+    return _copy_lazy_out_idx_kernel(n, block_n, _tl_dtype_str(dtype))
+
+
+@dataclass
+class LazyOutIdxCopyPrepared:
+    out: torch.Tensor
+    src: torch.Tensor
+    kernel: object
+
+    def run(self) -> None:
+        # out_idx makes TileLang allocate and return the output tensor. The
+        # training runtime keeps an out-variant API, so this adapter copies the
+        # result into the caller-owned output buffer.
+        self.out.copy_(self.kernel(self.src))
+
+    def destroy(self) -> None:
+        pass
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        self.destroy()
+
+
+def prepare_copy_lazy_out_idx(out: torch.Tensor, src: torch.Tensor) -> LazyOutIdxCopyPrepared:
+    if out.shape != src.shape:
+        raise ValueError("lazy out_idx copy expects matching shapes")
+    if out.dtype != src.dtype:
+        raise TypeError("lazy out_idx copy expects matching dtypes")
+    if not out.is_cuda or not src.is_cuda:
+        raise ValueError("lazy out_idx copy expects CUDA tensors")
+    if not out.is_contiguous() or not src.is_contiguous():
+        raise ValueError("lazy out_idx copy v1 supports contiguous tensors only")
+
+    n = src.numel()
+    block_n = _block_n(n)
+    if n % block_n != 0:
+        raise ValueError("lazy out_idx copy v1 requires N % BLOCK_N == 0")
+    kernel = _compiled_copy_lazy_out_idx(n, block_n, src.dtype)
+    return LazyOutIdxCopyPrepared(out, src, kernel)
+
+
+def copy_lazy_out_idx_(out: torch.Tensor, src: torch.Tensor) -> torch.Tensor:
+    with prepare_copy_lazy_out_idx(out, src) as prepared:
+        prepared.run()
+    return out
+
+
+def copy_lazy_out_idx(src: torch.Tensor) -> torch.Tensor:
+    n = src.numel()
+    block_n = _block_n(n)
+    if n % block_n != 0:
+        raise ValueError("lazy out_idx copy v1 requires N % BLOCK_N == 0")
+    kernel = _compiled_copy_lazy_out_idx(n, block_n, src.dtype)
+    return kernel(src)