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