def test_broadcast_binary_op(lhs_shape, rhs_shape, typ="add"):
global TASK
TASK = (
"bcast_binary_"
+ typ
+ "_lhs"
+ "_".join([str(ele) for ele in lhs_shape])
+ "rhs"
+ "_".join([str(ele) for ele in rhs_shape])
)
A = te.placeholder(shape=lhs_shape, name="A")
B = te.placeholder(shape=rhs_shape, name="B")
if typ == "add":
C = topi.broadcast_add(A, B)
elif typ == "sub":
C = topi.broadcast_sub(A, B)
elif typ == "div":
C = topi.broadcast_div(A, B)
elif typ == "mul":
C = topi.broadcast_mul(A, B)
elif typ == "maximum":
C = topi.broadcast_maximum(A, B)
elif typ == "minimum":
C = topi.broadcast_minimum(A, B)
else:
raise NotImplementedError
s = topi.cuda.schedule_broadcast(C)
fcuda = tvm.build(s, [A, B, C], "cuda", name="broadcast_binary" + "_" + typ)
lhs_npy = np.random.uniform(size=lhs_shape).astype(A.dtype)
rhs_npy = np.random.uniform(size=rhs_shape).astype(A.dtype)
if typ == "add":
out_npy = lhs_npy + rhs_npy
elif typ == "sub":
out_npy = lhs_npy - rhs_npy
elif typ == "div":
rhs_npy = np.abs(rhs_npy) + 0.001
out_npy = lhs_npy / rhs_npy
elif typ == "mul":
out_npy = lhs_npy * rhs_npy
elif typ == "maximum":
out_npy = np.maximum(lhs_npy, rhs_npy)
elif typ == "minimum":
out_npy = np.minimum(lhs_npy, rhs_npy)
lhs_nd = tvm.nd.array(lhs_npy, tvm.cuda())
rhs_nd = tvm.nd.array(rhs_npy, tvm.cuda())
out_nd = tvm.nd.array(np.empty(out_npy.shape).astype(B.dtype), tvm.cuda())
for _ in range(2):
fcuda(lhs_nd, rhs_nd, out_nd)
tvm.testing.assert_allclose(out_nd.numpy(), out_npy)
def mxnet_check():
"""This is a simple test function for MXNet bridge
It is not included as pytests, because of its dependency on mxnet
User can directly run this script to verify correctness.
"""
import mxnet as mx
from tvm import topi
import tvm
from tvm import te
import numpy as np
from tvm.contrib.mxnet import to_mxnet_func
# build a TVM function through topi
n = 20
shape = (20, )
scale = te.var("scale", dtype="float32")
x = te.placeholder(shape)
y = te.placeholder(shape)
z = topi.broadcast_add(x, y)
zz = te.compute(shape, lambda *i: z(*i) * scale)
target = tvm.target.cuda()
# build the function
with target:
s = topi.generic.schedule_injective(zz)
f = tvm.build(s, [x, y, zz, scale])
# get a mxnet version
mxf = to_mxnet_func(f, const_loc=[0, 1])
ctx = mx.gpu(0)
xx = mx.nd.uniform(shape=shape, ctx=ctx)
yy = mx.nd.uniform(shape=shape, ctx=ctx)
zz = mx.nd.empty(shape=shape, ctx=ctx)
# invoke myf: this runs in mxnet engine
mxf(xx, yy, zz, 10.0)
mxf(xx, yy, zz, 10.0)
tvm.testing.assert_allclose(zz.asnumpy(),
(xx.asnumpy() + yy.asnumpy()) * 10)