mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-20 06:30:07 +08:00
Lei Wang
0814b17129
* [Enhancement] Introduce non-restrict parameter support in code generation - Added a new PrimFunc-level attribute `tl.non_restrict_params` to specify handle Vars that should not be marked with the restrict qualifier during code generation. - Updated `CodeGenTileLangCPP`, `CodeGenTileLangCUDA`, and `CodeGenTileLangHIP` to handle non-restrict parameters, ensuring proper treatment of overlapping buffer aliases. - Implemented a new annotation function `annotate_restrict_buffers` to facilitate the marking of buffer parameters as non-restrict. - Enhanced the `SplitHostDevice` transformation to propagate non-restrict parameters from host to device functions. - Added a new transform function `HoistNonRestrictParams` to manage non-restrict parameters effectively. * [Enhancement] Improve HoistNonRestrictParams transformation - Updated the HoistNonRestrictParams function to recursively collect all `tl.non_restrict_params` annotations from nested blocks, enhancing flexibility in annotation placement. - Introduced a new NonRestrictCollector class to manage the collection and deduplication of non-restrict parameters. - Modified the SplitHostDevice transformation to remove the non-restrict attribute from the host-side PrimFunc after propagation to device kernels. - Adjusted the LowerAndLegalize function to directly apply the HoistNonRestrictParams transformation without exception handling, streamlining the process. * [Refactor] Simplify non-restrict parameter handling in code generation - Removed unnecessary normalization logic and associated data structures from `CodeGenTileLangCPP`, `CodeGenTileLangCUDA`, and `CodeGenTileLangHIP`. - Streamlined the handling of non-restrict parameters by directly inserting them into the `non_restrict` set, improving code clarity and maintainability. - Updated conditional checks to eliminate redundant checks against normalized names, enhancing performance and readability. * [Dependency] Update TVM subproject to latest commit 68aa8461 - Updated the TVM subproject to the latest commit, ensuring compatibility with recent changes and improvements. - Refactored non-restrict parameter handling in `CodeGenTileLangCPP`, `CodeGenTileLangCUDA`, and `CodeGenTileLangHIP` to enhance code clarity and maintainability. - Adjusted the `SplitHostDevice` transformation to streamline the propagation of non-restrict parameters. * fix
49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
import tilelang
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
|
|
|
|
def _get_sig_line(code: str) -> str:
|
|
# Find the kernel signature line in generated CUDA code
|
|
for line in code.splitlines():
|
|
line = line.strip()
|
|
if line.startswith('extern "C" __global__ void'):
|
|
return line
|
|
raise AssertionError("Kernel signature not found in generated code")
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_cuda_restrict_default_has_restrict():
|
|
N = 128
|
|
|
|
@T.prim_func
|
|
def kernel(x: T.Tensor((N,), T.float32), y: T.Tensor((N,), T.float32)):
|
|
with T.Kernel(N, threads=32) as pid:
|
|
y[pid] = x[pid] + 1.0
|
|
|
|
artifact = tilelang.lower(kernel, target="cuda")
|
|
sig = _get_sig_line(artifact.kernel_source)
|
|
# By default, kNoAlias is set and both pointers are restrict-qualified
|
|
assert "__restrict__" in sig
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_cuda_restrict_annotation_removes_restrict():
|
|
N = 128
|
|
|
|
@T.prim_func
|
|
def kernel_body_annot(x: T.Tensor((N,), T.float32), y: T.Tensor((N,), T.float32)):
|
|
# Explicitly mark buffers that may alias as non-restrict
|
|
with T.Kernel(N, threads=32) as pid:
|
|
T.annotate_restrict_buffers(x, y)
|
|
y[pid] = x[pid] + 1.0
|
|
|
|
art1 = tilelang.lower(kernel_body_annot, target="cuda")
|
|
sig1 = _get_sig_line(art1.kernel_source)
|
|
# No parameter should be emitted with __restrict__
|
|
assert "__restrict__" not in sig1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|