mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-21 06:40:23 +08:00
Lei Wang
00dd738835
* [Enhancement] Add read-only parameter annotation for CUDA codegen * Introduced the `AnnotateReadOnlyParams` transformation to annotate read-only handle parameters in PrimFuncs, enabling the generation of `const` qualifiers in CUDA codegen. * Updated `PrintFunctionSignature` and `AddFunction` methods to utilize the new attribute `tl.readonly_param_indices`, enhancing performance by allowing read-only cache loads. * Modified the optimization pipeline to include the new annotation step, improving the overall efficiency of the code generation process. * lint fix * [Dependency] Update apache-tvm-ffi version to >=0.1.3 * Updated the version of apache-tvm-ffi in pyproject.toml, requirements.txt, and requirements-dev.txt to ensure compatibility with the latest features and fixes. * Made adjustments in CUDA and HIP template files to use `const` qualifiers for global pointer parameters, enhancing code safety and clarity. * lint fix * [Enhancement] Refactor ReadWriteMarker for improved parameter handling * Updated the ReadWriteMarker class to accept a set of parameter or data variables, enhancing its ability to track written variables. * Introduced a new method, ResolveDataVarFromPtrArg, to resolve underlying buffer data from pointer-like arguments, improving accuracy in identifying written variables. * Modified the MarkReadOnlyParams function to gather handle parameters and their corresponding buffer data variables, streamlining the process of determining read-only parameters. * Enhanced the logic for identifying written variables to account for aliased data variables, ensuring comprehensive tracking of modifications. * lint fix * Update tma_load function to use const qualifier for global memory pointer * Changed the parameter type of gmem_ptr in the tma_load function from void* to void const* to enhance type safety and clarity in memory operations. * This modification ensures that the function correctly handles read-only global memory pointers, aligning with best practices in CUDA programming. * Remove commented-out code and reorder transformations in OptimizeForTarget function for clarity * Refactor buffer marking logic in annotate_read_only_params.cc to improve accuracy in identifying written variables. Update OptimizeForTarget function to reorder transformations for better clarity.
55 lines
2.0 KiB
Python
55 lines
2.0 KiB
Python
import tilelang.language as T
|
|
from tilelang.engine.lower import lower
|
|
from tilelang.jit.adapter.utils import match_declare_kernel
|
|
|
|
|
|
def _simple_add_kernel():
|
|
@T.prim_func
|
|
def main(
|
|
x: T.Tensor((128,), "float32"),
|
|
y: T.Tensor((128,), "float32"),
|
|
):
|
|
# One-dimensional kernel; writes y from x without modifying x
|
|
with T.Kernel(128, threads=32) as pid:
|
|
y[pid] = x[pid] + 1.0
|
|
|
|
return main
|
|
|
|
|
|
def test_codegen_emits_const_for_readonly_params():
|
|
# Lower without device compilation to retrieve CUDA source reliably
|
|
func = _simple_add_kernel()
|
|
artifact = lower(func, target="cuda", enable_device_compile=False)
|
|
|
|
src = artifact.kernel_source
|
|
print(src)
|
|
assert 'extern "C" __global__' in src
|
|
|
|
# Extract kernel signature and check qualifiers
|
|
lparen = match_declare_kernel(src)
|
|
rparen = src.find(")", lparen)
|
|
assert rparen != -1
|
|
signature = src[lparen:rparen]
|
|
|
|
# x is read-only: should be `const` and `__restrict__`
|
|
assert "const float* __restrict__" in signature
|
|
# y is written: must not be const, but still `__restrict__` due to noalias
|
|
# We ensure there is a non-const float* parameter with __restrict__ as well
|
|
assert "const float* __restrict__ x" in src or "const float *__restrict__ x" in src
|
|
assert " float* __restrict__ y" in src or " float *__restrict__ y" in src
|
|
|
|
# Also validate the function attribute carries read-only param indices
|
|
# Expect only the first handle parameter (x) to be marked read-only
|
|
device_mod = artifact.device_mod
|
|
prim_funcs = [f for f in device_mod.functions.values() if hasattr(f, "attrs")]
|
|
assert prim_funcs, "No PrimFunc found in device module"
|
|
pf = prim_funcs[0]
|
|
ro = pf.attrs.get("tl.readonly_param_indices")
|
|
assert ro is not None, "Expected tl.readonly_param_indices to be present"
|
|
ro_list = [int(i) for i in ro]
|
|
assert 0 in ro_list and 1 not in ro_list
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_codegen_emits_const_for_readonly_params()
|