mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-25 07:14:16 +08:00
KurisandLeiWang1999
d5d959ef7e
* [Fix] Add complement for non-var IterExpr * fix lint error * lint fix --------- Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
31 lines
847 B
Python
31 lines
847 B
Python
import tilelang as tl
|
|
import tilelang.testing
|
|
import tilelang.language as T
|
|
|
|
|
|
def test_issue_1374_non_var_itermark():
|
|
@tl.jit
|
|
def get_wrong_kernel(M: int = 4096):
|
|
dtype = "int32"
|
|
num_threads = 128
|
|
|
|
@T.prim_func
|
|
def main(A: T.Tensor((16, 14), dtype=dtype), B: T.Tensor((16, 448), dtype=dtype)):
|
|
with T.Kernel(1, threads=num_threads) as (bx,):
|
|
A_local = T.alloc_fragment((16, 14), dtype=dtype)
|
|
B_local = T.alloc_fragment((16, 448), dtype=dtype)
|
|
|
|
T.copy(A, A_local)
|
|
T.copy(B, B_local)
|
|
for i, j in T.Parallel(16, 448):
|
|
A_local[i, j // 32] += B[i, j]
|
|
|
|
return main
|
|
|
|
kernel = get_wrong_kernel()
|
|
print(kernel.get_kernel_source())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|