import tvm import numpy as np def enabled_ctx_list(): ctx_list = [('cpu', tvm.cpu(0)), ('gpu', tvm.gpu(0)), ('cl', tvm.opencl(0)), ('metal', tvm.metal(0)), ('rocm', tvm.rocm(0)), ('vulkan', tvm.vulkan(0)), ('vpi', tvm.vpi(0))] for k, v in ctx_list: assert tvm.context(k, 0) == v ctx_list = [x[1] for x in ctx_list if x[1].exist] return ctx_list ENABLED_CTX_LIST = enabled_ctx_list() print("Testing using contexts:", ENABLED_CTX_LIST) def test_nd_create(): for ctx in ENABLED_CTX_LIST: for dtype in ["uint8", "int8", "uint16", "int16", "uint32", "int32", "float32"]: x = np.random.randint(0, 10, size=(3, 4)) x = np.array(x, dtype=dtype) y = tvm.nd.array(x, ctx=ctx) z = y.copyto(ctx) assert y.dtype == x.dtype assert y.shape == x.shape assert isinstance(y, tvm.nd.NDArray) np.testing.assert_equal(x, y.asnumpy()) np.testing.assert_equal(x, z.asnumpy()) # no need here, just to test usablity ctx.sync() def test_fp16_conversion(): n = 100 for (src, dst) in [('float32', 'float16'), ('float16', 'float32')]: A = tvm.placeholder((n,), dtype=src) B = tvm.compute((n,), lambda i: A[i].astype(dst)) s = tvm.create_schedule([B.op]) func = tvm.build(s, [A, B], 'llvm') x_tvm = tvm.nd.array(100 * np.random.randn(n).astype(src) - 50) y_tvm = tvm.nd.array(100 * np.random.randn(n).astype(dst) - 50) func(x_tvm, y_tvm) expected = x_tvm.asnumpy().astype(dst) real = y_tvm.asnumpy() tvm.testing.assert_allclose(expected, real) if __name__ == "__main__": test_nd_create() test_fp16_conversion()