mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-22 06:50:25 +08:00
33 lines
785 B
Python
33 lines
785 B
Python
import tilelang
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
|
|
|
|
def test_var_assign() -> None:
|
|
|
|
@tilelang.jit(out_idx=-1)
|
|
def jit_kernel():
|
|
|
|
@T.prim_func
|
|
def test_var_assign(A: T.Tensor((2,), 'int32')):
|
|
with T.Kernel(1) as _:
|
|
a = T.alloc_var('int32', init=1)
|
|
b = T.alloc_var('int32', init=a) # b gets value of a
|
|
a = 2
|
|
d = T.alloc_var('int32', init=a) # c gets new value of a
|
|
A[0] = b
|
|
A[1] = d
|
|
|
|
print(test_var_assign)
|
|
return test_var_assign
|
|
|
|
kernel = jit_kernel()
|
|
print(kernel.get_kernel_source())
|
|
res = kernel()
|
|
assert res[0] == 1
|
|
assert res[1] == 2
|
|
|
|
|
|
if __name__ == '__main__':
|
|
tilelang.testing.main()
|