trending/tilelang
1
0
Fork 0
mirror of https://github.com/tile-ai/tilelang.git synced 2026-10-21 06:40:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Refactor] Improve assertion handling in CodeGenCHost and ArgBinder (#1352)

Browse Source 
* [Refactor] Improve assertion handling in CodeGenCHost and ArgBinder

This commit refines the assertion message generation in CodeGenCHost by optimizing the handling of equality checks and reducing buffer size for error messages. Additionally, it enhances the ArgBinder by introducing a nullable guard mechanism for assertions, allowing for more precise error handling when binding arguments. The changes improve the clarity and efficiency of assertion handling across the codebase.

* [Enhancement] Update matmul kernel and optimize argument binding

This commit enhances the matmul kernel by introducing additional tensor parameters and refining the pipeline stages for improved performance. It also updates the argument binding mechanism to include a flag indicating whether buffers are used, enhancing the efficiency of buffer management. Furthermore, the optimization phase in the engine is improved by adding a simplification step, ensuring better performance and clarity in the generated code.

* lint fix

* [Enhancement] Add tensor checks documentation and improve argument binding assertions

This commit introduces a new documentation page for host-side tensor checks, detailing the automatic validations performed by TileLang on kernel arguments. It enhances the ArgBinder by adding assertions for non-null pointers when arguments are used, improving error handling. Additionally, the optimization phase in the engine is updated to include a simplification step, ensuring better performance and clarity in the generated code.

* [Enhancement] Update .gitignore and refine matmul kernel for improved performance

This commit adds host checks logs to the .gitignore file to prevent unnecessary log files from being tracked. Additionally, it refines the matmul kernel by adjusting pipeline stages, updating tensor parameters, and enhancing argument handling for better performance. The changes also include improved error messages in the argument binding process, ensuring clearer diagnostics for users.

* lint fix

* lint fix

* [Refactor] Simplify tensor_null_test function and remove ptr_null_test

This commit refactors the tensor_null_test function by adding a with_bias parameter and removing the ptr_null_test function, which was previously unused. The run_test function is updated to reflect these changes, streamlining the testing process for tensor operations.

* lint fix

* fix
This commit is contained in:
Lei Wang
2025-11-28 03:28:14 +08:00
committed by GitHub
parent b8240b7ae9
commit 1e92d11cd2
27 changed files with 1101 additions and 284 deletions
Download Patch File Download Diff File Expand all files Collapse all files
testing/python/jit/test_tilelang_jit_nullptr.py
+12 -62
View File
@@ -7,49 +7,15 @@ from tilelang.utils import map_torch_type
@tl.jit
def ptr_null_test(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtype="float"):
@T.prim_func
def main(
a_ptr: T.ptr,
b_ptr: T.ptr,
c_ptr: T.ptr,
bias_ptr: T.ptr,
m: T.int32,
n: T.int32,
k: T.int32,
with_bias: T.bool,
):
A = T.make_tensor(a_ptr, (m, k), dtype)
B = T.make_tensor(b_ptr, (k, n), dtype)
C = T.make_tensor(c_ptr, (m, n), accum_dtype)
Bias = T.make_tensor(bias_ptr, (n), accum_dtype)
# Initialize Kernel Context
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
A_shared = T.alloc_shared((block_M, block_K), dtype)
B_shared = T.alloc_shared((block_N, block_K), dtype)
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
T.clear(C_local)
for ko in T.Pipelined(T.ceildiv(k, block_K), num_stages=3):
# Copy tile of A
T.copy(A[by * block_M, ko * block_K], A_shared)
T.copy(B[bx * block_N, ko * block_K], B_shared)
T.gemm(A_shared, B_shared, C_local, transpose_B=True)
if with_bias:
for i, j in T.Parallel(block_M, block_N):
C_local[i, j] += Bias[bx * block_N + j]
T.copy(C_local, C[by * block_M, bx * block_N])
return main
@tl.jit
def tensor_null_test(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtype="float"):
def tensor_null_test(M,
N,
K,
block_M,
block_N,
block_K,
dtype="float16",
accum_dtype="float",
with_bias=False):
@T.prim_func
def main(
@@ -57,7 +23,6 @@ def tensor_null_test(M, N, K, block_M, block_N, block_K, dtype="float16", accum_
B: T.Tensor((K, N), dtype),
C: T.Tensor((M, N), accum_dtype),
Bias: T.Tensor((N), accum_dtype),
with_bias: T.bool,
):
# Initialize Kernel Context
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
@@ -83,28 +48,13 @@ def tensor_null_test(M, N, K, block_M, block_N, block_K, dtype="float16", accum_
def run_test(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtype="float"):
kernel = ptr_null_test(M, N, K, block_M, block_N, block_K, dtype, accum_dtype)
a = torch.randn(M, K, device="cuda", dtype=map_torch_type(dtype))
b = torch.randn(N, K, device="cuda", dtype=map_torch_type(dtype))
c = torch.zeros(M, N, device="cuda", dtype=map_torch_type(accum_dtype))
d = torch.randn(N, device="cuda", dtype=map_torch_type(accum_dtype))
kernel(a, b, c, None, M, N, K, False)
ref_no_bias = (a @ b.T).to(map_torch_type(accum_dtype))
ref_with_bias = ref_no_bias + d
torch.testing.assert_close(c, ref_no_bias, atol=1e-2, rtol=1e-2)
kernel(a, b, c, d, M, N, K, True)
torch.testing.assert_close(c, ref_with_bias, atol=1e-2, rtol=1e-2)
kernel = tensor_null_test(M, N, K, block_M, block_N, block_K, dtype, accum_dtype)
kernel(a, b, c, None, False)
torch.testing.assert_close(c, ref_no_bias, atol=1e-2, rtol=1e-2)
kernel(a, b, c, d, True)
torch.testing.assert_close(c, ref_with_bias, atol=1e-2, rtol=1e-2)
kernel = tensor_null_test(
M, N, K, block_M, block_N, block_K, dtype, accum_dtype, with_bias=False)
kernel(a, b, c, None)
def test_nullptr():