def check(i, j, k): x = relay.var("x", shape=(i, k)) w1 = relay.var("w1", shape=(j, k)) w2 = relay.var("w2", shape=(2 * j, k)) w3 = relay.var("w3", shape=(3 * j, k)) y_before = before(x, w1, w2, w3) combine_pass = transform.CombineParallelDense(min_num_branches=3, to_batch=False) y = run_opt_pass(y_before, combine_pass) y_expected = expected(x, w1, w2, w3, j) y_expected = run_opt_pass(y_expected, transform.InferType()) tvm.ir.assert_structural_equal(y, y_expected, map_free_vars=True)
def check(i, j, k): x = relay.var("x", shape=(i, k)) w1 = relay.var("w1", shape=(j, k)) w2 = relay.var("w2", shape=(j, k)) w3 = relay.var("w3", shape=(j + 1, k)) w4 = relay.var("w4", shape=(j, k)) y_before = before(x, w1, w2, w3, w4) y = run_opt_pass(y_before, transform.CombineParallelDense(min_num_branches=2)) y_expected = expected(x, w1, w2, w3, w4) y_expected = run_opt_pass(y_expected, transform.InferType()) tvm.ir.assert_structural_equal(y, y_expected, map_free_vars=True)
def check(i, j, k): x = relay.var("x", shape=(i, k)) w1 = relay.var("w1", shape=(j, k)) w2 = relay.var("w2", shape=(j, k)) w3 = relay.var("w3", shape=(j + 1, k)) w4 = relay.var("w4", shape=(j, k)) y_before = before(x, w1, w2, w3, w4) y = run_opt_pass(y_before, transform.CombineParallelDense(min_num_branches=2)) y_expected = expected(x, w1, w2, w3, w4) y_expected = run_opt_pass(y_expected, transform.InferType()) assert relay.analysis.alpha_equal(y, y_expected)
def check(i, j, k, bias_shape1, bias_shape2): x = relay.var("x", shape=(i, k)) w1 = relay.var("w1", shape=(j, k)) w2 = relay.var("w2", shape=(2 * j, k)) b1 = relay.var("b1", shape=bias_shape1) b2 = relay.var("b2", shape=bias_shape2) y_before = before(x, w1, w2, b1, b2) combine_pass = transform.CombineParallelDense(min_num_branches=2, to_batch=False) y = run_opt_pass(y_before, combine_pass) y_expected = expected(x, w1, w2, b1, b2, j, bias_shape1, bias_shape2) y_expected = run_opt_pass(y_expected, transform.InferType()) tvm.ir.assert_structural_equal(y, y_expected, map_free_vars=True)
def check(i, j, k, scale1, scale2, newshape): x = relay.var("x", shape=(i, k)) w1 = relay.var("w1", shape=(j, k)) w2 = relay.var("w2", shape=(j, k)) b1 = relay.var("b1", shape=(j, )) b2 = relay.var("b2", shape=(j, )) scale1 = relay.var("scale1", shape=(1, )) scale2 = relay.var("scale2", shape=(1, )) y_before = before(x, w1, w2, b1, b2, scale1, scale2, newshape) y = run_opt_pass(y_before, transform.CombineParallelDense(min_num_branches=2)) y_expected = expected(x, w1, w2, b1, b2, scale1, scale2, newshape) y_expected = run_opt_pass(y_expected, transform.InferType()) tvm.ir.assert_structural_equal(y, y_expected, map_free_vars=True)
def check(i, j, k, scale1, scale2, newshape): x = relay.var("x", shape=(i, k)) w1 = relay.var("w1", shape=(j, k)) w2 = relay.var("w2", shape=(j, k)) b1 = relay.var("b1", shape=(j, )) b2 = relay.var("b2", shape=(j, )) scale1 = relay.var("scale1", shape=(1, )) scale2 = relay.var("scale2", shape=(1, )) y_before = before(x, w1, w2, b1, b2, scale1, scale2, newshape) y = run_opt_pass(y_before, transform.CombineParallelDense(min_num_branches=2)) y_expected = expected(x, w1, w2, b1, b2, scale1, scale2, newshape) y_expected = run_opt_pass(y_expected, transform.InferType()) assert relay.analysis.alpha_equal(y, y_expected)
def check(i, j, k, is_2d_bias): x = relay.var("x", shape=(i, k)) w1 = relay.var("w1", shape=(j, k)) w2 = relay.var("w2", shape=(j, k)) if is_2d_bias: b1 = relay.var("b1", shape=(i, j)) b2 = relay.var("b2", shape=(i, j)) else: b1 = relay.var("b1", shape=(j, )) b2 = relay.var("b2", shape=(j, )) y_before = before(x, w1, w2, b1, b2) y = run_opt_pass(y_before, transform.CombineParallelDense(min_num_branches=2)) y_expected = expected(x, w1, w2, b1, b2, is_2d_bias) y_expected = run_opt_pass(y_expected, transform.InferType()) tvm.ir.assert_structural_equal(y, y_expected, map_free_vars=True)
def check(i, j, k, is_2d_bias): x = relay.var("x", shape=(i, k)) w1 = relay.var("w1", shape=(j, k)) w2 = relay.var("w2", shape=(j, k)) if is_2d_bias: b1 = relay.var("b1", shape=(i, j)) b2 = relay.var("b2", shape=(i, j)) else: b1 = relay.var("b1", shape=(j, )) b2 = relay.var("b2", shape=(j, )) y_before = before(x, w1, w2, b1, b2) y = run_opt_pass(y_before, transform.CombineParallelDense(min_num_branches=2)) y_expected = expected(x, w1, w2, b1, b2, is_2d_bias) y_expected = run_opt_pass(y_expected, transform.InferType()) assert relay.analysis.alpha_equal(y, y_expected)
def run_combine_parallel(expr, min_num_branches=3): mod = tvm.IRModule.from_expr(expr) mod = transform.CombineParallelDense(min_num_branches)(mod) return mod["main"]