mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-22 06:50:25 +08:00
Lei Wang
f8ae600ce4
* remove debug print * Remove inline let expressions from the LowerAndLegalize function in phase.py * add test * Update sparse MLA examples to support SKV adjustment and correctness checks - Changed SKV parameter from 32768 to 8192 in sparse MLA backward and forward tests. - Added check_correctness parameter to test functions for validation of outputs. - Updated test cases to reflect new SKV values and correctness checks. * reduce test shape * Update documentation structure and refactor main function parameters in example_fusedmoe_tilelang.py - Added a new section for compiler internals in the documentation. - Refactored the main function in example_fusedmoe_tilelang.py to accept parameters for hidden dimensions, expert configurations, and batch/sequence sizes, improving flexibility and readability. * Update buffer access checks in merge_shared_memory_allocations.cc - Changed the condition for buffer access from less than (<) to less than or equal to (<=) to allow access at the same scope level. - Adjusted the logic for determining the access level when touching buffers to ensure correct handling of scope levels. * lint fix * Support pipeline with LetStmt * lint fix * • Fix LowerTileOp let handling to avoid LetInline dependency - inline let-bound BufferLoad nodes via resolver helpers and structured return - remap layouts/buffers using original data vars and only rewrite when needed - update pipeline planner to understand let-bound address_of buffers - document the new inline behaviour in docs/let_inline_fix.md * fix for wgmma pipeline with let binding * lint fix * test fix * reduce smem usage. * let binding enhancement * fix for dpgm * fix simplify * lint fix * use tilelang.Simplify instead of tir.Simplify * • Add TL_FORCE_LET_INLINE pass config and gate eager LetInline usage - register the new config in builtin headers/registration - add helper to pipeline enabling LetInline based on pass context - document LetStmt inlining controls and usage
52 lines
1.3 KiB
Python
52 lines
1.3 KiB
Python
import tilelang
|
|
import tilelang.testing
|
|
import tilelang.language as T
|
|
import torch
|
|
|
|
|
|
@tilelang.jit
|
|
def _tmp_var_kernel(N, block_N, dtype="float"):
|
|
|
|
@T.prim_func
|
|
def kernel(
|
|
A: T.Tensor((N,), dtype),
|
|
B: T.Tensor((N,), dtype),
|
|
):
|
|
with T.Kernel(T.ceildiv(N, block_N), threads=128) as bx:
|
|
for i in T.Parallel(block_N):
|
|
idx = bx * block_N + i
|
|
tmp = T.max(A[idx], 1)
|
|
B[idx] = tmp / 2
|
|
A[idx] = tmp * 2
|
|
|
|
return kernel
|
|
|
|
|
|
def run_tmp_var_test(N=1024, block_N=128):
|
|
kernel = _tmp_var_kernel(N, block_N)
|
|
|
|
a = torch.randn(N, device="cuda", dtype=torch.float)
|
|
b = torch.empty(N, device="cuda", dtype=torch.float)
|
|
|
|
a_ref = a.clone()
|
|
|
|
kernel(a, b)
|
|
|
|
# Reference computation
|
|
tmp_ref = torch.maximum(a_ref, torch.tensor(1.0, dtype=torch.float, device="cuda"))
|
|
b_ref = tmp_ref / 2
|
|
a_ref = tmp_ref * 2
|
|
|
|
# Validate correctness
|
|
tilelang.testing.torch_assert_close(a, a_ref, rtol=1e-2, atol=1e-2)
|
|
tilelang.testing.torch_assert_close(b, b_ref, rtol=1e-2, atol=1e-2)
|
|
|
|
|
|
def test_issue_814():
|
|
"""Test that temporary variables are correctly handled and not over-inlined"""
|
|
run_tmp_var_test(N=1024, block_N=128)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|