mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-23 07:00:29 +08:00
Kuris
5d806298de
* add verifier for parallel loop data racing * fix lint error * fix typo * fix lint error * fix bug in metal test
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):
|
|
B[i, j] += A_local[i, j // 32]
|
|
|
|
return main
|
|
|
|
kernel = get_wrong_kernel()
|
|
print(kernel.get_kernel_source())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|