mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-22 06:50:25 +08:00
KurisandLei Wang
bccb6485e4
* [Feature] Add support for A: T.Tensor(n + 1) and A: T.Tensor(2*n) * issue fix * fix * fix * decreate nproc for debugging --------- Co-authored-by: Lei Wang <leiwang1999@outlook.com>
72 lines
1.7 KiB
Python
72 lines
1.7 KiB
Python
import tilelang
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
import torch
|
|
|
|
|
|
def test_tensor_annot_mul():
|
|
|
|
@tilelang.jit
|
|
def example_tensor_annot():
|
|
n = T.symbolic('n')
|
|
|
|
@T.prim_func
|
|
def kernel(A: T.Tensor((n * 4,), T.int32),):
|
|
with T.Kernel(1) as _:
|
|
for i in range(n * 4):
|
|
A[i] = 0
|
|
|
|
return kernel
|
|
|
|
ker = example_tensor_annot()
|
|
A = torch.arange(16, dtype=torch.int32, device='cuda')
|
|
ker(A)
|
|
expected = torch.zeros(16, dtype=torch.int32, device='cuda')
|
|
assert torch.equal(A, expected)
|
|
|
|
|
|
def test_tensor_annot_add():
|
|
|
|
@tilelang.jit
|
|
def example_tensor_annot():
|
|
n = T.symbolic('n')
|
|
|
|
@T.prim_func
|
|
def kernel(A: T.Tensor((n + 1,), T.int32),):
|
|
with T.Kernel(1) as _:
|
|
for i in range(n + 1):
|
|
A[i] = 0
|
|
|
|
return kernel
|
|
|
|
ker = example_tensor_annot()
|
|
A = torch.arange(16, dtype=torch.int32, device='cuda')
|
|
ker(A)
|
|
expected = torch.zeros(16, dtype=torch.int32, device='cuda')
|
|
assert torch.equal(A, expected)
|
|
|
|
|
|
def test_tensor_annot_mul_add():
|
|
|
|
@tilelang.jit
|
|
def example_tensor_annot():
|
|
n = T.symbolic('n')
|
|
|
|
@T.prim_func
|
|
def kernel(A: T.Tensor((n * 3 + 1,), T.int32),):
|
|
with T.Kernel(1) as _:
|
|
for i in range(n * 3 + 1):
|
|
A[i] = 0
|
|
|
|
return kernel
|
|
|
|
ker = example_tensor_annot()
|
|
A = torch.arange(16, dtype=torch.int32, device='cuda')
|
|
ker(A)
|
|
expected = torch.zeros(16, dtype=torch.int32, device='cuda')
|
|
assert torch.equal(A, expected)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
tilelang.testing.main()
|