()
| 22 | |
| 23 | |
| 24 | def test_dltensor_compatible(): |
| 25 | @I.ir_module |
| 26 | class Module: |
| 27 | @T.prim_func(s_tir=True) |
| 28 | def arange(A: T.handle): |
| 29 | n = T.int32() |
| 30 | Ab = T.match_buffer(A, (n,), "int64") |
| 31 | for i in T.serial(n - 1): |
| 32 | Ab[i + 1] = Ab[i] + T.int64(1) |
| 33 | |
| 34 | mod = Module |
| 35 | f = tvm.compile(mod, target="llvm") |
| 36 | a = tvm.runtime.tensor(np.zeros(10, dtype="int64")) |
| 37 | f(a) |
| 38 | np.testing.assert_equal(a.numpy(), np.arange(a.shape[0])) |
| 39 | |
| 40 | |
| 41 | if __name__ == "__main__": |