def test_llvm_persist_parallel():
n = 128
A = te.placeholder((n, ), name="A")
B = te.compute(A.shape, lambda *i: A(*i) + 1, name="B")
C = te.compute(A.shape, lambda *i: te.sqrt(B(*i)) * 2 + 2, name="C")
s = te.create_schedule(C.op)
xo, xi = s[C].split(C.op.axis[0], factor=8)
xo1, xo2 = s[C].split(xo, nparts=1)
s[B].compute_at(s[C], xo1)
s[B].parallel(s[B].op.axis[0])
s[B].pragma(s[B].op.axis[0], "parallel_barrier_when_finish")
s[C].parallel(xi)
s[C].pragma(xo1, "parallel_launch_point")
s[C].pragma(xi, "parallel_stride_pattern")
def check_llvm():
# BUILD and invoke the kernel.
f = tvm.build(s, [A, C], "llvm")
dev = tvm.cpu(0)
# launch the kernel.
a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev)
c = tvm.nd.array(np.zeros(n, dtype=C.dtype), dev)
f(a, c)
tvm.testing.assert_allclose(c.numpy(),
np.sqrt(a.numpy() + 1) * 2 + 2,
rtol=1e-5)
check_llvm()
def sqrt(x):
"""Take square root of input x.
Parameters
----------
x : tvm.te.Tensor
Input argument.
Returns
-------
y : tvm.te.Tensor
The result.
"""
return te.compute(x.shape, lambda *i: te.sqrt(x(*i)))
def norm_bmn( # pylint: disable=invalid-name,missing-docstring
B: int,
M: int,
N: int,
) -> Tuple[te.Tensor, te.Tensor]:
a = te.placeholder((B, M, N), name="A")
i = te.reduce_axis((0, M), name="i")
j = te.reduce_axis((0, N), name="j")
c = te.compute(
(B, ),
lambda b: te.sum(a[b][i][j] * a[b][i][j], axis=[i, j]),
name="C",
)
d = te.compute((B, ), lambda b: te.sqrt(c[b]), name="D")
return (a, d)
def test_basic_operation():
np.random.seed(0)
shape = (10, 10)
x = te.var("x", dtype='float32')
k = te.reduce_axis((0, 10), name="k")
l = te.reduce_axis((0, 10), name="l")
A0 = te.placeholder(shape, name='A0')
A1 = te.placeholder(shape, name='A1')
zeros = np.zeros(shape)
B = te.compute(shape, lambda i, j: A0[i, j], name='B')
check_grad(B, [A0])
B = te.compute(shape, lambda i, j: A0[i, j] + A1[i, j], name='B')
check_grad(B, [A0, A1])
B = te.compute(shape, lambda i, j: A0[i, j] + A0[j, i], name='B')
check_grad(B, A0)
B = te.compute(shape, lambda i, j: te.floor(A0[i, j]), name='B')
check_grad(B, A0, desired_grads=[zeros])
B = te.compute(shape, lambda i, j: te.ceil(A0[i, j]), name='B')
check_grad(B, A0, desired_grads=[zeros])
B = te.compute(shape, lambda i, j: te.trunc(A0[i, j]), name='B')
check_grad(B, A0, desired_grads=[zeros])
B = te.compute(shape, lambda i, j: te.round(A0[i, j]), name='B')
check_grad(B, A0, desired_grads=[zeros])
B = te.compute(shape, lambda i, j: A0[i, j] + te.exp(A0[j, i]), name='B')
check_grad(B, A0)
B = te.compute(
shape,
lambda i, j: te.log(0.1 + te.abs(A0[i, j] + te.exp(A0[j, i]))),
name='B')
check_grad(B, A0)
B = te.compute(shape,
lambda i, j: te.sigmoid(A0[i, j] * A0[i, j] * A0[j, i]),
name='B')
check_grad(B, A0)
B = te.compute(shape,
lambda i, j: te.tanh(A0[i, j] * A0[i, j] * A0[j, i]),
name='B')
check_grad(B, A0)
B = te.compute(shape,
lambda i, j: te.sqrt(A0[i, j] * A0[i, j] * A0[j, i]),
name='B')
check_grad(B, A0, data_range=(0.1, 10))
B = te.compute(shape,
lambda i, j: te.power(te.abs(A0[i, j]), A0[j, i]),
name='B')
check_grad(B, A0, data_range=(-4, 4))
B = te.compute(shape, lambda i, j: A0[i, j] * A0[j, i], name='B')
check_grad(B, A0)
B = te.compute((10, ),
lambda i: te.sum(A0[i, k] * A0[k, i], axis=k),
name='B')
check_grad(B, A0)
B = te.compute(shape,
lambda i, j: te.sum(A0[i, k] * A0[k, i] + 5, axis=k),
name='B')
check_grad(B, A0)
B = te.compute(shape,
lambda i, j: te.max(A0[i, k] * A0[k, j] + 5, axis=k),
name='B')
check_grad(B, A0)
B = te.compute(shape,
lambda i, j: A0[i, j] * (A1[j, i] + A0[j, i]),
name='B')
check_grad(B, [A0, A1])
B = te.compute(shape,
lambda i, j: te.sum(
A0[k, k] - A0[te.min(j + k, 9), j] * A0[i, k], axis=k),
name='B')
check_grad(B, A0)
def fcombine(x, y):
return x * y
def fidentity(t0):
return tvm.tir.const(1, t0)
prod = te.comm_reducer(fcombine, fidentity, name='prod')
B = te.compute((10, 10),
lambda i, j: prod(A0[i, k] + A0[k, i], axis=k),
name='B')
check_grad(B, A0)
X = te.placeholder((10, ), name='X')
A = te.compute((10, ), lambda i: X[i] + X[9 - i])
B = te.compute((10, ), lambda i: X[i] * X[9 - i])
Y = topi.tensordot(A, B, 1)
check_grad(Y, X)
