def check_vulkan(dtype, n): if not tvm.vulkan(0).exist or not tvm.runtime.enabled("vulkan"): print("skip because vulkan is not enabled..") return A = tvm.placeholder((n, ), name='A', dtype=dtype) ctx = tvm.vulkan(0) a_np = np.random.uniform(size=(n, )).astype(A.dtype) a = tvm.nd.empty((n, ), A.dtype, ctx).copyfrom(a_np) b_np = a.asnumpy() tvm.testing.assert_allclose(a_np, b_np) tvm.testing.assert_allclose(a_np, a.asnumpy())
def worker(): A = te.placeholder((n,), name="A", dtype="float32") B = te.placeholder((n,), name="B", dtype="float32") functions = [ ( lambda: te.compute((n,), lambda i: 2 * A[i] + 3 * B[i]), lambda a, b: 2 * a + 3 * b, ), (lambda: te.compute((n,), lambda i: A[i] + B[i]), lambda a, b: a + b), (lambda: te.compute((n,), lambda i: A[i] + 2 * B[i]), lambda a, b: a + 2 * b), ] def build_f(f_ref): (C_f, ref) = f_ref C = C_f() s = te.create_schedule(C.op) xo, xi = s[C].split(C.op.axis[0], factor=num_thread) s[C].bind(xo, bx) s[C].bind(xi, tx) fun = tvm.build(s, [A, B, C], "vulkan") return (fun, ref) fs = [ build_f(random.choice(functions)) for _ in range(np.random.randint(low=1, high=10)) ] dev = tvm.vulkan(0) a = tvm.nd.empty((n,), A.dtype, dev).copyfrom(np.random.uniform(size=(n,))) b = tvm.nd.empty((n,), B.dtype, dev).copyfrom(np.random.uniform(size=(n,))) cs = [tvm.nd.empty((n,), A.dtype, dev) for _ in fs] for ((f, _), c) in zip(fs, cs): f(a, b, c) for ((_, ref), c) in zip(fs, cs): tvm.testing.assert_allclose(c.numpy(), ref(a.numpy(), b.numpy()))
def check_vulkan(dtype, n): A = te.placeholder((n, ), name="A", dtype=dtype) dev = tvm.vulkan(0) a_np = np.random.uniform(size=(n, )).astype(A.dtype) a = tvm.nd.empty((n, ), A.dtype, dev).copyfrom(a_np) b_np = a.numpy() tvm.testing.assert_allclose(a_np, b_np) tvm.testing.assert_allclose(a_np, a.numpy())
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
def check_vulkan(dtype, n, lanes): if not tvm.vulkan(0).exist or not tvm.runtime.enabled("vulkan"): print("skip because vulkan is not enabled..") return A = tvm.placeholder((n, ), name='A', dtype="%sx%d" % (dtype, lanes)) B = tvm.compute((n, ), lambda i: A[i] + tvm.const(1, A.dtype), name='B') s = tvm.create_schedule(B.op) xo, xi = s[B].split(B.op.axis[0], factor=num_thread) s[B].bind(xo, bx) s[B].bind(xi, tx) fun = tvm.build(s, [A, B], "vulkan") ctx = tvm.vulkan(0) a = tvm.nd.empty((n, ), A.dtype, ctx).copyfrom(np.random.uniform(size=(n, lanes))) c = tvm.nd.empty((n, ), B.dtype, ctx) fun(a, c) tvm.testing.assert_allclose(c.asnumpy(), a.asnumpy() + 1)
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
def check_vulkan(dtype, n, lanes): A = te.placeholder((n,), name="A", dtype="%sx%d" % (dtype, lanes)) B = te.compute((n,), lambda i: A[i] + tvm.tir.const(1, A.dtype), name="B") s = te.create_schedule(B.op) xo, xi = s[B].split(B.op.axis[0], factor=num_thread) s[B].bind(xo, bx) s[B].bind(xi, tx) fun = tvm.build(s, [A, B], "vulkan") dev = tvm.vulkan(0) a = tvm.nd.empty((n,), A.dtype, dev).copyfrom(np.random.uniform(size=(n, lanes))) c = tvm.nd.empty((n,), B.dtype, dev) fun(a, c) tvm.testing.assert_allclose(c.numpy(), a.numpy() + 1)
def requires_gpu(*args): """Mark a test as requiring a GPU to run. Tests with this mark will not be run unless a gpu is present. Parameters ---------- f : function Function to mark """ _requires_gpu = [ pytest.mark.skipif( not tvm.cuda().exist and not tvm.rocm().exist and not tvm.opencl().exist and not tvm.metal().exist and not tvm.vulkan().exist, reason="No GPU present", ), *uses_gpu(), ] return _compose(args, _requires_gpu)
def test_scalar_params(num_int_params): n = te.var("n") scalars = [te.var("scale{}".format(i)) for i in range(num_int_params)] scalar_sum = scalars[0] for s in scalars[1:]: scalar_sum += s A = te.placeholder((n, ), name="A") B = te.compute(A.shape, lambda i: scalar_sum + A[i], name="B") s = te.create_schedule(B.op) xo, xi = s[B].split(B.op.axis[0], factor=64) s[B].bind(xo, bx) s[B].bind(xi, tx) f_add = tvm.build(s, scalars + [A, B], target) n = 1024 scalars = [1 for _ in scalars] dev = tvm.vulkan(0) a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev) b = tvm.nd.array(np.zeros(n, dtype=B.dtype), dev) f_add(*scalars, a, b) tvm.testing.assert_allclose(a.numpy() + sum(scalars), b.numpy())
###################################################################### # Execute the portable graph on TVM # --------------------------------- # Now, we would like to reproduce the same forward computation using TVM. from tvm.contrib import graph_runtime if target == 'llvm': ctx = tvm.cpu(0) elif target == 'cuda': ctx = tvm.gpu(0) elif target == 'opengl': ctx = tvm.opengl(0) elif target == 'opencl': ctx = tvm.cl(0) elif target == 'vulkan': ctx = tvm.vulkan(0) elif target == 'metal': ctx = tvm.metal(0) else: raise ValueError('No supported context type for ' % target) dtype = 'float32' m = graph_runtime.create(graph, lib, ctx) # set inputs m.set_input('data', tvm.nd.array(x.astype(dtype))) m.set_input(**params) # execute m.run() # get outputs tvm_output = m.get_output(0) top1 = np.argmax(tvm_output.asnumpy()[0])