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()