mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-22 06:50:25 +08:00
Lei WangandZhiwen Mo
f5d9da4678
* Refactor GEMM and Reduce operations by moving NormalizeToBufferRegion and MakeAccessPtrFromRegion to utils.{h,cc} for better code organization and reuse.
* lint fix
* Refactor region handling by removing the RegionOp and updating NormalizeToBufferRegion to only accept BufferLoad and BufferRegion. This change improves code organization and simplifies the handling of memory regions across various operations.
* fix
* Refactor memory region handling by introducing `tl.region` calls across various operations, including GEMM and fill functions. This change enhances the consistency of region management and improves code organization by utilizing utility functions for buffer region conversions.
* fix
* fix
* test fix
* lint fix
* Refactor GEMM operations to improve memory region handling by replacing `mbarPtr_` with `mbarRegion_` and updating related logic in both C++ and Python implementations. This change enhances the clarity and consistency of buffer region management.
* fix
* lint fix
* fix
* fix
* test fix
* lint fix
* lint fix
* minor fix
* fix
---------
Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk>
82 lines
2.1 KiB
Python
82 lines
2.1 KiB
Python
# ruff: noqa
|
|
|
|
import torch
|
|
import tilelang
|
|
import tilelang.testing
|
|
import tilelang.language as T
|
|
|
|
|
|
@tilelang.jit
|
|
def _empty_kernel():
|
|
|
|
@T.prim_func
|
|
def empty_kernel():
|
|
with T.Kernel(1, threads=32) as thread_idx:
|
|
pass
|
|
|
|
return empty_kernel
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_empty_kernel_lowering():
|
|
# Ensure a valid CUDA runtime context is current on this thread for the
|
|
# target device before using driver API calls. Without this, calls like
|
|
# cuModuleLoadData can fail with CUDA_ERROR_INVALID_CONTEXT, especially
|
|
# for kernels that don't touch any device memory or streams beforehand
|
|
# (e.g., "empty" kernels) and therefore haven't triggered context
|
|
# creation implicitly.
|
|
torch.cuda.set_device(0)
|
|
kernel = _empty_kernel()
|
|
kernel()
|
|
|
|
|
|
@tilelang.jit
|
|
def _empty_with_dead_code_kernel():
|
|
num_tokens = T.dynamic("num_tokens")
|
|
|
|
@T.prim_func
|
|
def buggy_kernel(x: T.Tensor[(num_tokens,), "float32"]):
|
|
with T.Kernel(num_tokens, threads=32) as pid:
|
|
y = x[pid]
|
|
|
|
return buggy_kernel
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_empty_with_dead_code_kernel():
|
|
kernel = _empty_with_dead_code_kernel()
|
|
x = torch.randn((128,), dtype=torch.float32, device="cuda")
|
|
kernel(x)
|
|
|
|
|
|
@tilelang.jit
|
|
def _empty_kernel_with_binding_variants(use_tuple_binding: bool = False):
|
|
|
|
@T.prim_func
|
|
def kernel_with_tuple_kernel_binding():
|
|
with T.Kernel(1, threads=32) as (pid,):
|
|
print(pid)
|
|
pass
|
|
|
|
@T.prim_func
|
|
def kernel_with_scalar_kernel_binding():
|
|
with T.Kernel(1, threads=32) as pid:
|
|
print(pid)
|
|
pass
|
|
|
|
return kernel_with_tuple_kernel_binding if use_tuple_binding else kernel_with_scalar_kernel_binding
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_empty_kernel_with_binding_variants():
|
|
torch.cuda.set_device(0)
|
|
kernel = _empty_kernel_with_binding_variants()
|
|
kernel()
|
|
|
|
tuple_kernel = _empty_kernel_with_binding_variants(use_tuple_binding=True)
|
|
tuple_kernel()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|