def check_expr_sum(dtype): n = 4 a = tvm.placeholder((n,n,n), name="a", dtype=dtype) b = tvm.placeholder((n,n,n), name="b", dtype=dtype) c = tvm.compute(a.shape, lambda i, j, k: tvm.trace([a[i][j][k]],"tvm.trace_callback3") + tvm.trace([b[i][j][k]],"tvm.trace_callback3")) s = tvm.create_schedule(c.op) f = tvm.build(s, [a, b, c]) xnd = tvm.nd.array(np.array(np.ones((n,n,n), dtype=a.dtype))) ynd = tvm.nd.array(np.array(np.ones((n,n,n), dtype=b.dtype))) znd = tvm.nd.array(np.zeros((n,n,n), dtype=c.dtype)) f(xnd, ynd, znd) assert(np.array_equal(znd.asnumpy(), xnd.asnumpy() + ynd.asnumpy()))
def check_expr_sum(dtype): n = 4 a = tvm.placeholder((n, n, n), name="a", dtype=dtype) b = tvm.placeholder((n, n, n), name="b", dtype=dtype) c = tvm.compute( a.shape, lambda i, j, k: tvm.trace([a[i][j][k]], "tvm.trace_callback3") + tvm.trace([b[i][j][k]], "tvm.trace_callback3")) s = tvm.create_schedule(c.op) f = tvm.build(s, [a, b, c]) xnd = tvm.nd.array(np.array(np.ones((n, n, n), dtype=a.dtype))) ynd = tvm.nd.array(np.array(np.ones((n, n, n), dtype=b.dtype))) znd = tvm.nd.array(np.zeros((n, n, n), dtype=c.dtype)) f(xnd, ynd, znd) assert (np.array_equal(znd.asnumpy(), xnd.asnumpy() + ynd.asnumpy()))
def check_expr_sum_custom(dtype): n = 4 a = tvm.placeholder((n,n), name="a", dtype=dtype) b = tvm.placeholder((n,n), name="b", dtype=dtype) c = tvm.compute(a.shape, lambda i,j: tvm.trace([a[i][j]], "tvm.trace_callback4") + tvm.trace([b[i][j]], "tvm.trace_callback4")) s = tvm.create_schedule(c.op) f = tvm.build(s, [a, b, c]) npa = np.array([[1,0,0,0], [0,1,0,0],[0,0,1,0],[0,0,0,1]], dtype=a.dtype) npb = np.array([[1,0,0,0], [0,1,0,0],[0,0,1,0],[0,0,0,1]], dtype=a.dtype) xnd = tvm.nd.array(npa) ynd = tvm.nd.array(npb) znd = tvm.nd.array(np.zeros((n,n), dtype=c.dtype)) f(xnd, ynd, znd) assert(np.array_equal(znd.asnumpy(), npa + npb))
def check_assign(dtype): n = 4 x = tvm.placeholder((n,), name="X", dtype=dtype) y = tvm.compute(x.shape, lambda i: tvm.trace([x[i]], "tvm.trace_change_float_first")) z = tvm.compute(x.shape, lambda i: tvm.trace([y[i]], "tvm.trace_change_float_second")) s = tvm.create_schedule(z.op) f = tvm.build(s, [x, y, z], "llvm") xnd = tvm.nd.array(np.ones((n,), dtype=x.dtype)) ynd = tvm.nd.array(np.zeros((n,), dtype=y.dtype)) znd = tvm.nd.array(np.zeros((n,), dtype=z.dtype)) f(xnd, ynd, znd) check_array_first = np.array([13.0, 13.0, 13.0, 13.0]) check_array_second = np.array([14.0, 14.0, 14.0, 14.0]) assert(np.array_equal(ynd.asnumpy(), check_array_first)) assert(np.array_equal(znd.asnumpy(), check_array_second))
def check_assign(dtype): n = 4 x = tvm.placeholder((n,n,n), name="X", dtype=dtype) y = tvm.compute(x.shape, lambda i, j, k: tvm.trace([x[i][j][k]], "tvm.trace_callback2")) z = tvm.compute(x.shape, lambda i, j, k: tvm.trace([y[i][j][k]], "tvm.trace_callback2")) s = tvm.create_schedule(z.op) f = tvm.build(s, [x, y, z], "llvm") xnd = tvm.nd.array(np.ones((n,n,n), dtype=x.dtype)) ynd = tvm.nd.array(np.zeros((n,n,n), dtype=y.dtype)) znd = tvm.nd.array(np.zeros((n,n,n), dtype=z.dtype)) f(xnd, ynd, znd) assert(np.array_equal(xnd.asnumpy(), np.ones((n,n,n)))) assert(np.array_equal(ynd.asnumpy(), np.ones((n,n,n)))) assert(np.array_equal(znd.asnumpy(), np.ones((n,n,n))))
def test_trace_default_action(): n = 2 x = tvm.placeholder((n, n, n), name="X", dtype="float32") y = tvm.compute(x.shape, lambda i, j, k: tvm.trace([i, j, k, x[i][j][k]])) s = tvm.create_schedule(y.op) f = tvm.build(s, [x, y], target="llvm") xnd = tvm.nd.array(np.ones((n, n, n), dtype=x.dtype)) ynd = tvm.nd.array(np.zeros((n, n, n), dtype=y.dtype)) f(xnd, ynd)
def test_trace_default_action(): n = 2 x = tvm.placeholder((n,n,n), name="X", dtype="float32") y = tvm.compute(x.shape, lambda i, j, k: tvm.trace([i, j, k, x[i][j][k]])) s = tvm.create_schedule(y.op) f = tvm.build(s, [x, y], target="llvm") xnd = tvm.nd.array(np.ones((n,n,n), dtype=x.dtype)) ynd = tvm.nd.array(np.zeros((n,n,n), dtype=y.dtype)) f(xnd, ynd)
def check_assign(dtype): n = 4 x = tvm.placeholder((n, ), name="X", dtype=dtype) y = tvm.compute( x.shape, lambda i: tvm.trace([x[i]], "tvm.trace_change_int_first")) z = tvm.compute( x.shape, lambda i: tvm.trace([y[i]], "tvm.trace_change_int_second")) s = tvm.create_schedule(z.op) f = tvm.build(s, [x, y, z], "llvm") xnd = tvm.nd.array(np.ones((n, ), dtype=x.dtype)) ynd = tvm.nd.array(np.zeros((n, ), dtype=y.dtype)) znd = tvm.nd.array(np.zeros((n, ), dtype=z.dtype)) f(xnd, ynd, znd) check_array_first = np.array([13, 13, 13, 13]) check_array_second = np.array([14, 14, 14, 14]) assert (np.array_equal(ynd.asnumpy(), check_array_first)) assert (np.array_equal(znd.asnumpy(), check_array_second))
def check_expr_sum_custom(dtype): n = 4 a = tvm.placeholder((n, n), name="a", dtype=dtype) b = tvm.placeholder((n, n), name="b", dtype=dtype) c = tvm.compute( a.shape, lambda i, j: tvm.trace([a[i][j]], "tvm.trace_callback4") + tvm.trace([b[i][j]], "tvm.trace_callback4")) s = tvm.create_schedule(c.op) f = tvm.build(s, [a, b, c]) npa = np.array( [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]], dtype=a.dtype) npb = np.array( [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]], dtype=a.dtype) xnd = tvm.nd.array(npa) ynd = tvm.nd.array(npb) znd = tvm.nd.array(np.zeros((n, n), dtype=c.dtype)) f(xnd, ynd, znd) assert (np.array_equal(znd.asnumpy(), npa + npb))
def check_assign(dtype): n = 4 x = tvm.placeholder((n, n, n), name="X", dtype=dtype) y = tvm.compute( x.shape, lambda i, j, k: tvm.trace([x[i][j][k]], "tvm.trace_callback2")) z = tvm.compute( x.shape, lambda i, j, k: tvm.trace([y[i][j][k]], "tvm.trace_callback2")) s = tvm.create_schedule(z.op) f = tvm.build(s, [x, y, z], "llvm") xnd = tvm.nd.array(np.ones((n, n, n), dtype=x.dtype)) ynd = tvm.nd.array(np.zeros((n, n, n), dtype=y.dtype)) znd = tvm.nd.array(np.zeros((n, n, n), dtype=z.dtype)) f(xnd, ynd, znd) assert (np.array_equal(xnd.asnumpy(), np.ones((n, n, n)))) assert (np.array_equal(ynd.asnumpy(), np.ones((n, n, n)))) assert (np.array_equal(znd.asnumpy(), np.ones((n, n, n))))
def trace2(): s = (4, ) x = tvm.placeholder(s, name="x", dtype="int32") y = tvm.compute(x.shape, lambda i: tvm.trace([x[i]]) + 100) # print(tvm.lower(s, [x, y], simple_mode=True)) s = tvm.create_schedule(y.op) f = tvm.build(s, [x, y]) xnd = tvm.nd.array(2 * np.ones((4, ), dtype=x.dtype)) ynd = tvm.nd.array(np.zeros((4, ), dtype=y.dtype)) f(xnd, ynd) print(ynd)
def check_expr_sum(dtype): n = 4 a = tvm.placeholder((n,n,n), name="a", dtype=dtype) b = tvm.placeholder((n,n,n), name="b", dtype=dtype) e = tvm.placeholder((n,n,n), name="e", dtype=dtype) d = tvm.placeholder((n,n,n), name="d", dtype=dtype) c = tvm.compute(a.shape, lambda i, j, k: tvm.trace([i, j, k, a[i][j][k]], "tvm.trace_silent") + tvm.trace([i, j, k, b[i][j][k]], "tvm.trace_silent") + tvm.trace([i, j, k, d[i][j][k]], "tvm.trace_silent") + tvm.trace([i, j, k, e[i][j][k]], "tvm.trace_silent")) s = tvm.create_schedule(c.op) f = tvm.build(s, [a, b, d, e, c]) a_nd = tvm.nd.array(np.array(np.ones((n,n,n), dtype=a.dtype))) b_nd = tvm.nd.array(np.array(np.ones((n,n,n), dtype=b.dtype))) d_nd = tvm.nd.array(np.array(np.ones((n,n,n), dtype=d.dtype))) e_nd = tvm.nd.array(np.array(np.ones((n,n,n), dtype=e.dtype))) c_nd = tvm.nd.array(np.zeros((n,n,n), dtype=c.dtype)) f(a_nd, b_nd, d_nd, e_nd, c_nd) assert(np.array_equal(c_nd.asnumpy(), a_nd.asnumpy() + b_nd.asnumpy() + d_nd.asnumpy() + e_nd.asnumpy()))
def check_expr_sum(dtype): n = 4 a = tvm.placeholder((n, n, n), name="a", dtype=dtype) b = tvm.placeholder((n, n, n), name="b", dtype=dtype) e = tvm.placeholder((n, n, n), name="e", dtype=dtype) d = tvm.placeholder((n, n, n), name="d", dtype=dtype) c = tvm.compute( a.shape, lambda i, j, k: tvm.trace([i, j, k, a[i][j][ k]], "tvm.trace_silent") + tvm.trace([i, j, k, b[i][j][ k]], "tvm.trace_silent") + tvm.trace([i, j, k, d[i][j][ k]], "tvm.trace_silent") + tvm.trace( [i, j, k, e[i][j][k]], "tvm.trace_silent")) s = tvm.create_schedule(c.op) f = tvm.build(s, [a, b, d, e, c]) a_nd = tvm.nd.array(np.array(np.ones((n, n, n), dtype=a.dtype))) b_nd = tvm.nd.array(np.array(np.ones((n, n, n), dtype=b.dtype))) d_nd = tvm.nd.array(np.array(np.ones((n, n, n), dtype=d.dtype))) e_nd = tvm.nd.array(np.array(np.ones((n, n, n), dtype=e.dtype))) c_nd = tvm.nd.array(np.zeros((n, n, n), dtype=c.dtype)) f(a_nd, b_nd, d_nd, e_nd, c_nd) assert (np.array_equal( c_nd.asnumpy(), a_nd.asnumpy() + b_nd.asnumpy() + d_nd.asnumpy() + e_nd.asnumpy()))