diff --git a/examples/05_tilelang_copy_modes.py b/examples/05_tilelang_copy_modes.py new file mode 100644 index 0000000..adf4f6d --- /dev/null +++ b/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() diff --git a/ops/common/tilelang/eager_copy.py b/ops/common/tilelang/eager_copy.py new file mode 100644 index 0000000..c26d227 --- /dev/null +++ b/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) diff --git a/ops/common/tilelang/lazy_out_idx_copy.py b/ops/common/tilelang/lazy_out_idx_copy.py new file mode 100644 index 0000000..49b136f --- /dev/null +++ b/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)