"""Test code for broadcasting operators.""" from common import get_all_backend import numpy as np import tvm import topi def verify_broadcast_to_ele(in_shape, out_shape, fbcast): # Build the logic and compile the function A = tvm.placeholder(shape=in_shape, name="A") B = fbcast(A, out_shape) def check_device(device): ctx = tvm.context(device, 0) if not ctx.exist: print("Skip because %s is not enabled" % device) return print("Running on target: %s" % device) with tvm.target.create(device): s = topi.generic.schedule_broadcast(B) foo = tvm.build(s, [A, B], device, name="broadcast_to") data_npy = np.random.uniform(size=in_shape).astype(A.dtype) out_npy = np.broadcast_to(data_npy, out_shape) data_nd = tvm.nd.array(data_npy, ctx) out_nd = tvm.nd.array(np.empty(out_shape).astype(B.dtype), ctx) foo(data_nd, out_nd) tvm.testing.assert_allclose(out_nd.asnumpy(), out_npy) for target in get_all_backend(): check_device(target) check_device("sdaccel") def verify_broadcast_binary_ele(lhs_shape, rhs_shape, ftopi, fnumpy, lhs_min=-100, lhs_max=100, rhs_min=-100, rhs_max=100, dtype="float32"): # Build the logic and compile the function A = (tvm.var("A", dtype=dtype) if lhs_shape is None else tvm.placeholder(shape=lhs_shape, name="A", dtype=dtype)) B = (tvm.var("B", dtype=dtype) if rhs_shape is None else tvm.placeholder(shape=rhs_shape, name="B", dtype=dtype)) C = ftopi(A, B) if isinstance(A, tvm.expr.Expr) and isinstance(B, tvm.expr.Expr): assert(isinstance(C, tvm.expr.Expr)) return def check_device(device): ctx = tvm.context(device, 0) if not ctx.exist: print("Skip because %s is not enabled" % device) return print("Running on target: %s" % device) with tvm.target.create(device): s = topi.generic.schedule_broadcast(C) foo = tvm.build(s, [A, B, C], device, name="broadcast_binary" + "_" + ftopi.__name__) if lhs_shape is None: lhs_npy = float(np.random.uniform(low=lhs_min, high=lhs_max)) if dtype.startswith('int'): lhs_npy = int(lhs_npy) lhs_nd = lhs_npy else: lhs_npy = np.random.uniform(low=lhs_min, high=lhs_max, size=lhs_shape).astype(A.dtype) lhs_nd = tvm.nd.array(lhs_npy, ctx) if rhs_shape is None: rhs_npy = float(np.random.uniform(low=rhs_min, high=rhs_max)) if dtype.startswith('int'): rhs_npy = int(rhs_npy) rhs_nd = rhs_npy else: rhs_npy = np.random.uniform(low=rhs_min, high=rhs_max, size=rhs_shape).astype(A.dtype) rhs_nd = tvm.nd.array(rhs_npy, ctx) out_npy = fnumpy(lhs_npy, rhs_npy) out_nd = tvm.nd.array(np.empty(out_npy.shape).astype(C.dtype), ctx) foo(lhs_nd, rhs_nd, out_nd) tvm.testing.assert_allclose(out_nd.asnumpy(), out_npy, rtol=1E-4, atol=1E-4) for target in get_all_backend(): check_device(target) check_device("sdaccel") def test_broadcast_to(): verify_broadcast_to_ele((1,), (10,), topi.broadcast_to) verify_broadcast_to_ele((), (10,), topi.broadcast_to) verify_broadcast_to_ele((1, 1, 5, 4), (3, 4, 4, 4, 5, 4), topi.broadcast_to) verify_broadcast_to_ele((1, 128, 1, 32), (64, 128, 64, 32), topi.broadcast_to) def test_add(): verify_broadcast_binary_ele( (), (), topi.add, np.add) verify_broadcast_binary_ele( (5, 2, 3), (2, 1), topi.add, np.add) def test_subtract(): verify_broadcast_binary_ele( (5, 2, 3), (), topi.subtract, np.subtract) verify_broadcast_binary_ele( (5, 2, 3), None, topi.subtract, np.subtract) verify_broadcast_binary_ele( None, None, topi.subtract, np.subtract) verify_broadcast_binary_ele( (1, 32), (64, 32), topi.subtract, np.subtract) def test_multiply(): verify_broadcast_binary_ele( (5, 64, 128), (2, 5, 64, 1), topi.multiply, np.multiply) def test_divide(): verify_broadcast_binary_ele( None, (10,), topi.divide, np.divide, rhs_min=0.0001) verify_broadcast_binary_ele( (), None, topi.divide, np.divide, rhs_min=0.0001) verify_broadcast_binary_ele( (2, 3, 1, 32), (64, 32), topi.divide, np.divide, rhs_min=0.0001) def test_maximum_minmum(): verify_broadcast_binary_ele( (32,), (64, 32), topi.maximum, np.maximum) verify_broadcast_binary_ele( (1, 2, 2, 1, 32), (64, 32), topi.minimum, np.minimum) def test_power(): verify_broadcast_binary_ele( (1, 2, 2), (2,), topi.power, np.power, lhs_min=0.001, rhs_min=0.001, rhs_max=2) def test_mod(): verify_broadcast_binary_ele( (1, 2, 2), (2,), topi.mod, np.mod, lhs_min=0.001, rhs_min=1, dtype="int32") def test_cmp(): # explicit specify the output type def greater(x, y): return topi.greater(x, y).astype("int8") def less(x, y): return topi.less(x, y).astype("int8") def equal(x, y): return topi.equal(x, y).astype("int8") def not_equal(x, y): return topi.not_equal(x, y).astype("int8") def greater_equal(x, y): return topi.greater_equal(x, y).astype("int8") def less_equal(x, y): return topi.less_equal(x, y).astype("int8") verify_broadcast_binary_ele( (1, 2, 2), (2,), greater, np.greater) verify_broadcast_binary_ele( (2, 1, 2), (2, 3, 1), less, np.less) verify_broadcast_binary_ele( (2, 1, 2), (2, 3, 1), equal, np.equal, lhs_min=-2, lhs_max=2, rhs_min=-2, rhs_max=2, dtype='int32') verify_broadcast_binary_ele( (2, 1, 2), (2, 3, 1), not_equal, np.not_equal, lhs_min=-2, lhs_max=2, rhs_min=-2, rhs_max=2, dtype='int32') verify_broadcast_binary_ele( (7, 1, 5), (7, 3, 1), greater_equal, np.greater_equal, lhs_min=-3, lhs_max=3, rhs_min=-3, rhs_max=3, dtype='int32') verify_broadcast_binary_ele( (7, 1, 5), (7, 3, 1), less_equal, np.less_equal, lhs_min=-3, lhs_max=3, rhs_min=-3, rhs_max=3, dtype='int32') def test_shift(): # explicit specify the output type verify_broadcast_binary_ele( (2, 1, 2), None, topi.right_shift, np.right_shift, dtype="int32", rhs_min=0, rhs_max=32) verify_broadcast_binary_ele( (1, 2, 2), (2,), topi.left_shift, np.left_shift, dtype="int32", rhs_min=0, rhs_max=32) verify_broadcast_binary_ele( (1, 2, 2), (2,), topi.left_shift, np.left_shift, dtype="int8", rhs_min=0, rhs_max=32) if __name__ == "__main__": test_add() test_shift() test_cmp() test_mod() test_subtract() test_multiply() test_divide() test_maximum_minmum() test_power() test_broadcast_to()