def test_bound_vars(): x = relay.Var("x") y = relay.Var("y") z = relay.Var("z") a = relay.Var("a") f1 = relay.Function([x, y, z], relay.Let(a, x, relay.Tuple([]))) assert_vars_match(bound_vars(f1), [x, y, z, a]) tup = relay.Tuple([x, y, z, a]) assert len(bound_vars(tup)) == 0 f2 = relay.Function([x, y], relay.Tuple([x, y, z, a])) assert_vars_match(bound_vars(f2), [x, y])
def expected(): x = relay.var("x", shape=(), dtype="bool") y = relay.var("y", shape=(), dtype="float32") cond_t = relay.const(True) cond_f = relay.const(False) one = relay.const(1, dtype="float32") two = relay.const(2, dtype="float32") three = relay.const(3, dtype="float32") y2 = relay.var("y2") true_branch = relay.If(cond_t, y2, relay.add(three, y2)) true_branch = relay.Let(y2, relay.add(y, y), true_branch) false_branch = relay.If(cond_f, two, one) body = relay.If(x, true_branch, false_branch) return body
def test_let(): lv = relay.Var('x') ty = None arr = tvm.nd.array(10) value = relay.Constant(arr) # I would prefer that the order of arguments # matches syntax let x: t = v in b let = relay.Let(lv, value, lv) assert let.var == lv assert let.value == value assert let.body == lv assert let.span == None str(let) check_json_roundtrip(let)
def test_ref(): shape = (10, 10) dtype = 'float32' t = relay.TensorType(shape, dtype) x = relay.var("x", t) r = relay.Var("r") u = relay.Var("u") body = relay.RefRead(r) body = relay.Let(u, relay.RefWrite(r, relay.RefRead(r) + relay.RefRead(r)), body) body = relay.Let(r, relay.RefCreate(x), body) func = relay.Function([x], body) func = run_infer_type(func) back_func = run_infer_type(gradient(func)) assert back_func.checked_type == relay.FuncType( [t], relay.TupleType([t, relay.TupleType([t])])) x_nd = rand(dtype, *shape) ex = create_executor() forward, (grad_x, ) = ex.evaluate(back_func)(x_nd) tvm.testing.assert_allclose(forward.asnumpy(), 2 * x_nd.asnumpy()) tvm.testing.assert_allclose(grad_x.asnumpy(), 2 * np.ones_like(grad_x.asnumpy()))
def test_order(): z = relay.const(3) y = relay.const(2) x = relay.const(1) val = x + y * z check_eval(val, 7.0) anf = transform.OptimizeOnExpr( val, [transform.ToANormalForm(), transform.InferType()]) a = relay.Var('a', relay.IncompleteType()) b = relay.Var('b', relay.IncompleteType()) c = relay.Var('c', relay.IncompleteType()) d = relay.Var('d', relay.IncompleteType()) e = relay.Var('e', relay.IncompleteType()) expected_output = e expected_output = relay.Let(e, a + d, expected_output) expected_output = relay.Let(d, b * c, expected_output) expected_output = relay.Let(c, z, expected_output) expected_output = relay.Let(b, y, expected_output) expected_output = relay.Let(a, x, expected_output) expected_output = transform.OptimizeOnExpr(expected_output, transform.InferType()) assert alpha_equal(anf, expected_output)
def test_match(): # pair each match keyword with whether it specifies a complete match or not match_keywords = [("match", True), ("match?", False)] for (match_keyword, is_complete) in match_keywords: mod = tvm.IRModule() list_var = relay.GlobalTypeVar("List") typ_var = relay.TypeVar("A") cons_constructor = relay.Constructor( "Cons", [typ_var, list_var(typ_var)], list_var) nil_constructor = relay.Constructor("Nil", [], list_var) list_def = relay.TypeData(list_var, [typ_var], [cons_constructor, nil_constructor]) mod[list_var] = list_def length_var = relay.GlobalVar("length") typ_var = relay.TypeVar("A") input_type = list_var(typ_var) input_var = relay.Var("xs", input_type) rest_var = relay.Var("rest") cons_case = relay.Let( relay.var("", type_annotation=None), UNIT, relay.add(relay.const(1), relay.Call(length_var, [rest_var]))) body = relay.Match(input_var, [ relay.Clause( relay.PatternConstructor( cons_constructor, [relay.PatternWildcard(), relay.PatternVar(rest_var)]), cons_case), relay.Clause(relay.PatternConstructor(nil_constructor, []), relay.const(0)) ], complete=is_complete) length_func = relay.Function([input_var], body, int32, [typ_var]) mod[length_var] = length_func assert_parse_module_as( """ %s def @length[A](%%xs: List[A]) -> int32 { %s (%%xs) { Cons(_, %%rest : List[A]) => { (); 1 + @length(%%rest) }, Nil => 0, } } """ % (LIST_DEFN, match_keyword), mod)
def test_function_invalidate(): shape = () dtype = "bool" t = relay.TensorType(shape, dtype) d = relay.Var("d", t) r = relay.Var("r") fetch = relay.Function([], relay.RefRead(r)) fet = relay.Var("fetch") fet_obscured = relay.Var("fetch_obscured") u = relay.Var("u") body = relay.If(d, fet_obscured(), fet_obscured()) body = relay.Let(u, relay.RefWrite(r, relay.const(1)), body) body = relay.Let(fet_obscured, relay.If(d, fet, fet), body) body = relay.Let(fet, fetch, body) body = relay.Let(r, relay.RefCreate(relay.const(0)), body) f = relay.Function([d], body) f = infer_type(f) pe_f = infer_type(partial_evaluate(f)) ex = create_executor() f_res = ex.evaluate(f)(relay.const(True)) pe_f_res = ex.evaluate(pe_f)(relay.const(True)) np.testing.assert_allclose(f_res.asnumpy(), np.ones_like(f_res.asnumpy())) np.testing.assert_allclose(pe_f_res.asnumpy(), np.ones_like(pe_f_res.asnumpy()))
def test_ref_execution_order(): # we want to have effects execute from left to right x = relay.Var('x') y = relay.Var('y') f = relay.Var('f') r = relay.Var('r') expr = relay.Let( f, relay.Function([x, y], x), # r = 1 relay.Let( r, relay.RefCreate(relay.const(1)), relay.Tuple([ # should be 1 relay.RefRead(r), # set r to 2 and read back seq(relay.RefWrite(r, relay.const(2)), relay.RefRead(r)), # set r to 3 and read back seq(relay.RefWrite(r, relay.const(3)), relay.RefRead(r)), # set r to 4 and read as first arg to f # set r to 5 and read as second arg to f # f should evaluate to 4 f(seq(relay.RefWrite(r, relay.const(4)), relay.RefRead(r)), seq(relay.RefWrite(r, relay.const(5)), relay.RefRead(r))), # read back 5 relay.RefRead(r) ]))) tup_val = run_as_python(expr) assert_adt_len(tup_val, 5) assert_tensor_value(tup_val[0], 1) assert_tensor_value(tup_val[1], 2) assert_tensor_value(tup_val[2], 3) assert_tensor_value(tup_val[3], 4) assert_tensor_value(tup_val[4], 5)
def test_match_pattern(): mod = tvm.IRModule() box, box_ctor = init_box_adt(mod) v = relay.Var("v") w = relay.Var("w") match = relay.Let( v, box_ctor(relay.const(1)), relay.Match(v, [ relay.Clause( relay.PatternConstructor(box_ctor, [relay.PatternVar(w)]), w) ]), ) match_val = run_as_python(match, mod) assert_tensor_value(match_val, 1)
def test_arbitrary_let_nesting(): # something that is tricky to do in Python but comes naturally in Relay mod = tvm.IRModule() p = Prelude(mod) x = relay.Var("x") r = relay.Var("r") y = relay.Var("y") z = relay.Var("z") expr = relay.Tuple([ relay.Let(x, relay.Tuple([relay.const(1), relay.const(2)]), relay.TupleGetItem(x, 1)), relay.Let( r, relay.RefCreate(relay.const(1)), seq(relay.RefWrite(r, relay.const(3)), relay.RefRead(r)), ), relay.Let(y, p.id(relay.Let(z, relay.const(4), z)), y), ]) tup_val = run_as_python(expr, mod) assert_adt_len(tup_val, 3) assert_tensor_value(tup_val[0], 2) assert_tensor_value(tup_val[1], 3) assert_tensor_value(tup_val[2], 4)
def test_tuple_type(): assert alpha_equal( relay.fromtext( """ let %_: () = (); () """), relay.Let( relay.Var("_", relay.TupleType([])), UNIT, UNIT ) ) assert alpha_equal( relay.fromtext( """ let %_: (int32,) = (0,); () """), relay.Let( relay.Var("_", relay.TupleType([int32])), relay.Tuple([relay.const(0)]), UNIT ) ) assert alpha_equal( relay.fromtext( """ let %_: (int32, int32) = (0, 1); () """), relay.Let( relay.Var("_", relay.TupleType([int32, int32])), relay.Tuple([relay.const(0), relay.const(1)]), UNIT ) )
def test_tensor_type(): assert_parses_as( "let %_ : Tensor[(), float32] = (); ()", relay.Let(relay.Var("_", relay.TensorType((), "float32")), UNIT, UNIT), ) assert_parses_as( "let %_ : Tensor[(1), float32] = (); ()", relay.Let(relay.Var("_", relay.TensorType((1, ), "float32")), UNIT, UNIT), ) assert_parses_as( "let %_ : Tensor[(1, 1), float32] = (); ()", relay.Let(relay.Var("_", relay.TensorType((1, 1), "float32")), UNIT, UNIT), ) assert_parses_as( "let %_ : Tensor[(?, 1), float32] = (); ()", relay.Let( relay.Var("_", relay.TensorType((tvm.tir.Any(), 1), "float32")), UNIT, UNIT), )
def test_tuple_alpha_equal(): v1 = relay.Var("v1") v2 = relay.Var("v2") # unit value is a valid tuple assert alpha_equal(relay.Tuple([]), relay.Tuple([])) tup = relay.Tuple([v1, relay.const(2), relay.const(3), relay.Tuple([relay.const(4)])]) same = relay.Tuple([v1, relay.const(2), relay.const(3), relay.Tuple([relay.const(4)])]) assert alpha_equal(tup, same) # use the eq_map let_tup = relay.Let(v1, tup, v1) let_mapped = relay.Let(v2, relay.Tuple([v2, relay.const(2), relay.const(3), relay.Tuple([relay.const(4)])]), v2) assert alpha_equal(let_tup, let_mapped) more_fields = relay.Tuple([v1, relay.const(2), relay.const(3), relay.Tuple([relay.const(4)]), v2]) assert not alpha_equal(tup, more_fields) fewer_fields = relay.Tuple([v1, relay.const(2), relay.const(3)]) assert not alpha_equal(tup, fewer_fields) different_end = relay.Tuple([v1, relay.const(2), relay.const(3), relay.Tuple([relay.const(5)])]) assert not alpha_equal(tup, different_end) different_start = relay.Tuple([v2, relay.const(2), relay.const(3), relay.Tuple([relay.const(4)])]) assert not alpha_equal(tup, different_start) longer_at_end = relay.Tuple([v1, relay.const(2), relay.const(3), relay.Tuple([relay.const(4), relay.const(5)])]) assert not alpha_equal(tup, longer_at_end)
def test_nested_match_pattern(): mod = relay.Module() box, box_ctor = init_box_adt(mod) v = relay.Var('v') w = relay.Var('w') match = relay.Let( v, box_ctor(box_ctor(relay.const(2))), relay.Match(v, [ relay.Clause( relay.PatternConstructor(box_ctor, [ relay.PatternConstructor(box_ctor, [relay.PatternVar(w)]) ]), w) ])) match_val = run_as_python(match, mod) assert_tensor_value(match_val, 2)
def test_let_if_scope(): x = relay.var("x", "float32") y = relay.var("y", "float32") cond = relay.var("cond", "bool") sb = relay.ScopeBuilder() with sb.if_scope(cond): v1 = sb.let("v", relay.const(1, "float32")) v2 = sb.let("v", x) sb.ret(relay.subtract(v1, v2)) with sb.else_scope(): v3 = relay.var("v") let2 = relay.Let(v3, y, v3) sb.ret(relay.add(let2, let2)) result = sb.get() f = relay.Function([x, y, cond], result) print(f.astext())
def test_all_vars(): x = relay.Var("x") y = relay.Var("y") z = relay.Var("z") f1 = relay.Function([x, y], z) assert_vars_match(all_vars(f1), [x, y, z]) f2 = relay.Function([x], relay.Let(y, relay.Tuple([]), z)) assert_vars_match(all_vars(f2), [x, y, z]) f3 = relay.Function([x], relay.Tuple([y, z])) assert_vars_match(all_vars(f3), [x, y, z]) tup = relay.Tuple([x, y, z]) assert_vars_match(all_vars(tup), [x, y, z])
def _test_tuple(mode): shape = (10, 10) dtype = "float32" t = relay.TensorType(shape, dtype) x = relay.var("x", t) y = relay.var("y", t) z = relay.var("z", t) if mode == "higher_order": tup = relay.Var("tup") func = relay.Function( [x, y, z], relay.Let( tup, relay.Tuple([x, y, z]), relay.TupleGetItem(tup, 0) + relay.TupleGetItem(tup, 1) - relay.TupleGetItem(tup, 2), ), ) else: # first order does not do let. tup = relay.Tuple([x, y, z]) func = relay.Function( [x, y, z], relay.TupleGetItem(tup, 0) + relay.TupleGetItem(tup, 1) - relay.TupleGetItem(tup, 2), ) func = run_infer_type(func) back_func = run_infer_type(gradient(func, mode=mode)) assert back_func.checked_type == relay.FuncType( [t, t, t], relay.TupleType([t, relay.TupleType([t, t, t])])) x_nd = rand(dtype, *shape) y_nd = rand(dtype, *shape) z_nd = rand(dtype, *shape) x_np = x_nd.asnumpy() y_np = y_nd.asnumpy() z_np = z_nd.asnumpy() expected_forward = x_np + y_np - z_np ex = create_executor() forward, (grad_x, grad_y, grad_z) = ex.evaluate(back_func)(x_nd, y_nd, z_nd) tvm.testing.assert_allclose(forward.asnumpy(), expected_forward) tvm.testing.assert_allclose(grad_x.asnumpy(), np.ones_like(grad_x.asnumpy())) tvm.testing.assert_allclose(grad_y.asnumpy(), np.ones_like(grad_y.asnumpy())) tvm.testing.assert_allclose(grad_z.asnumpy(), -1 * np.ones_like(grad_z.asnumpy()))
def test_match_order(): mod = relay.Module() box, box_ctor = init_box_adt(mod) v = relay.Var('v') w = relay.Var('w') # wildcard pattern goes first match = relay.Let( v, box_ctor(box_ctor(relay.const(2))), relay.Match(v, [ relay.Clause(relay.PatternWildcard(), relay.const(1)), relay.Clause( relay.PatternConstructor(box_ctor, [ relay.PatternConstructor(box_ctor, [relay.PatternVar(w)]) ]), w) ])) match_val = run_as_python(match, mod) assert_tensor_value(match_val, 1)
def before(): var1 = relay.var("var1", shape=(2,)) var2 = relay.var("var2", shape=(), dtype="int32") var3 = relay.var("var3", shape=(2,)) cond = relay.less(var2, relay.const(10, dtype="int32")) loop = relay.var("while_loop") ii = var2 + relay.const(1, dtype="int32") ss = var3 + var1 true_branch = loop(ii, ss) ife = relay.If(cond, true_branch, var3) func_1 = relay.Function([var2, var3], ife) ret = relay.Let(loop, func_1, loop(relay.const(0, dtype="int32"), relay.zeros_like(var1))) func_2 = relay.Function([var1], ret) mod = tvm.IRModule.from_expr(func_2) return mod
def before(x, conv_weight, in_bias, in_scale, channels): args = [x, conv_weight, in_bias] x = relay.multiply(x, in_scale) x = relay.nn.relu(x) x = relay.add(x, in_bias) x_var = relay.Var("x_var") y1 = relay.nn.conv2d( x_var, conv_weight, channels=channels, kernel_size=(3, 3), data_layout="NHWC", kernel_layout="HWIO", padding=(1, 1), ) z = relay.add(y1, x) let = relay.Let(x_var, x, z) return relay.Function(args, let)
def after(): var1 = relay.var("var1", shape=(2, )) var2 = relay.var("var2", shape=(), dtype="int32") var3 = relay.var("var3", shape=(2, )) var4 = relay.const(10, dtype="int32") cb_1 = relay.annotation.compiler_begin(var2, target) cb_2 = relay.annotation.compiler_begin(var4, target) less_condition = relay.less(cb_1, cb_2) ce_1 = relay.annotation.compiler_end(less_condition, target) loop = relay.var("while_loop") # if condition cb_3 = relay.annotation.compiler_begin(var2, target) cb_4 = relay.annotation.compiler_begin(relay.const(1, dtype="int32"), target) add_op_1 = relay.add(cb_3, cb_4) ce_2 = relay.annotation.compiler_end(add_op_1, target) cb_5 = relay.annotation.compiler_begin(ce_2, "default") cb_6 = relay.annotation.compiler_begin(var3, target) cb_7 = relay.annotation.compiler_begin(var1, target) add_op_2 = relay.add(cb_6, cb_7) ce_3 = relay.annotation.compiler_end(add_op_2, target) cb_8 = relay.annotation.compiler_begin(ce_3, "default") true_branch = loop(cb_5, cb_8) # while loop ce_4 = relay.annotation.compiler_end(true_branch, "default") if_condition = relay.If(ce_1, ce_4, var3) cb_9 = relay.annotation.compiler_begin(relay.const(0, dtype="int32"), "default") cb_10 = relay.annotation.compiler_begin(var1, target) zeros_like = relay.zeros_like(cb_10) ce_5 = relay.annotation.compiler_end(zeros_like, target) cb_11 = relay.annotation.compiler_begin(ce_5, "default") while_condition = loop(cb_9, cb_11) ce_6 = relay.annotation.compiler_end(while_condition, "default") func_1 = relay.Function([var2, var3], if_condition) ret = relay.Let(loop, func_1, ce_6) func_2 = relay.Function([var1], ret) mod = tvm.IRModule.from_expr(func_2) return mod
def expected_if_expr(x): """ free_var %x: float32 let %v1: float32 = add(%x, 1f /* ty=float32 */) /* ty=float32 */; %0 = equal(%x, 2f /* ty=float32 */) /* ty=bool */; if (%0) { multiply(%v1, 2f /* ty=float32 */) /* ty=float32 */ } else { multiply(%v1, 1f /* ty=float32 */) /* ty=float32 */ } """ one = relay.const(1, dtype="float32") two = relay.const(2, dtype="float32") v1 = relay.var("v1") v2 = relay.equal(x, two) true_branch = relay.multiply(v1, two) false_branch = relay.multiply(v1, one) body = relay.If(v2, true_branch, false_branch) body = relay.Let(v1, relay.add(x, one), body) return body
def test_local_recursion(): mod = tvm.IRModule() p = Prelude(mod) _, cons, nil = p.mod.get_type("List") v = relay.Var("v") h = relay.Var("h") t = relay.Var("t") f = relay.Var("f") # just returns the same list let = relay.Let( f, relay.Function( [v], relay.Match( v, [ relay.Clause( relay.PatternConstructor( cons, [relay.PatternVar(h), relay.PatternVar(t)]), cons(h, f(t)), ), relay.Clause(relay.PatternConstructor(nil, []), nil()), ], ), ), f( cons(relay.const(1), cons(relay.const(2), cons(relay.const(3), nil())))), ) val = run_as_python(let, mod) assert_constructor_value(val, cons, 2) assert_tensor_value(val.fields[0], 1) assert_constructor_value(val.fields[1], cons, 2) assert_tensor_value(val.fields[1].fields[0], 2) assert_constructor_value(val.fields[1].fields[1], cons, 2) assert_tensor_value(val.fields[1].fields[1].fields[0], 3) assert_constructor_value(val.fields[1].fields[1].fields[1], nil, 0)
def test_head_cons(): mod = relay.Module() p = Prelude(mod) def hd_impl(): a = relay.TypeVar("a") x = relay.Var("x", p.l(a)) y = relay.Var("y") z = relay.Var("z") cons_case = relay.Clause(relay.PatternConstructor(p.cons, [relay.PatternVar(y), relay.PatternVar(z)]), y) return relay.Function([x], relay.Match(x, [cons_case]), a, [a]) t = relay.TypeVar("t") x = relay.Var("x", t) hd = relay.Var("hd") body = relay.Let(hd, hd_impl(), hd(p.cons(x, p.nil()))) f = relay.Function([x], body, None, [t]) f = infer_type(f, mod=mod) res = dcpe(f) assert alpha_equal(res, relay.Function([x], x, t, [t]))
def test_match_effect_exactly_once(): mod = relay.Module() p = Prelude(mod) # the list should be of length 1! # Unless we mistakenly execute the data clause more than once r = relay.Var('r') data = seq(relay.RefWrite(r, p.cons(relay.Tuple([]), relay.RefRead(r))), relay.RefRead(r)) match = relay.Let( r, relay.RefCreate(p.nil()), relay.Match(data, [ relay.Clause(relay.PatternConstructor(p.nil, []), relay.const(0)), relay.Clause( relay.PatternConstructor( p.cons, [relay.PatternWildcard(), relay.PatternConstructor(p.nil, [])]), relay.const(1)), relay.Clause(relay.PatternWildcard(), relay.const(2)) ])) match_val = run_as_python(match, mod) assert_tensor_value(match_val, 1)
def convert_func(self, graph): """Convert a graph.""" for p in graph.parameters: self.node_map[p] = self.on_parameter(p) params = [self.ref(p) for p in graph.parameters] seq = [] for node in toposort(graph.output, NodeVisitor(), in_graph(graph)): if node in self.node_map: continue elif node.is_constant_graph() and node.value.parent is None: self.node_map[node] = self.graph_map[node.value] else: self.node_map[node] = relay.var(f"seq.{self.i}") self.i += 1 seq.append(node) out = self.ref(graph.output) for op in reversed(seq): var = self.node_map[op] if op.is_apply(): val = self.on_apply(op) elif op.is_constant_graph(): val = self.convert_func(op.value) elif op.is_constant(): val = self.on_constant(op) # This forces the rebuild of constants every time they # are encountered since they may be shared amongst # multiple graphs and it causes problems otherwise. del self.node_map[op] else: raise AssertionError(f"Bad node for sequence: {op}") out = relay.Let(var, val, out) return relay.Function(params, out, ret_type=to_relay_type(graph.output.abstract))
def test_valid_if(): cond = relay.var("cond", dtype="bool", shape=()) shared = relay.var("shared") true_branch = shared false_branch = relay.add(shared, shared) body = relay.If(cond, true_branch, false_branch) shared_bound = relay.var("shared_bound", shape=(1, ), dtype="float32") body = relay.Let(shared, shared_bound, body) """ The program below uses let binding to control the scope of %shared, which follows the basic block normal form. free_var %shared_bound: Tensor[(1), float32] let %shared = %shared_bound; free_var %cond: bool if (%cond) { %shared } else { add(%shared, %shared) } """ check_basic_block_normal_form(body)
def test_recursive(): mod = tvm.IRModule() x = relay.var('x', shape=(2,)) i = relay.var('i', shape=(), dtype='int32') s = relay.var('s', shape=(2,)) cond = i < relay.const(10, dtype='int32') loop = relay.var('while_loop') sb = relay.scope_builder.ScopeBuilder() with sb.if_scope(cond): ii = i + relay.const(1, dtype='int32') ss = s + x sb.ret(loop(ii, ss)) with sb.else_scope(): sb.ret(s) func = relay.Function([i, s], sb.get()) ret = relay.Let(loop, func, loop(relay.const(0, dtype='int32'), relay.zeros(shape=(2,), dtype='float32'))) mod["main"] = relay.Function([x], ret) new_mod = transform.LambdaLift()(mod) assert len(new_mod.functions) == 2
def test_valid_if2(): """ fn (%x: float32) { let %v1 = add(%x, 1f); %0 = equal(%x, 2f); if (%0) { multiply(%v1, 2f) } else { multiply(%v1, 1f) } } """ x = relay.var("x", shape=(), dtype="float32") one = relay.const(1, dtype="float32") two = relay.const(2, dtype="float32") v1 = relay.var("v1") v2 = relay.equal(x, two) true_branch = relay.multiply(v1, two) false_branch = relay.multiply(v1, one) body = relay.If(v2, true_branch, false_branch) body = relay.Let(v1, relay.add(x, one), body) func = relay.Function([x], body) check_basic_block_normal_form(func)
def seq(*exprs): ret = exprs[0] for expr in exprs[1:]: ret = relay.Let(relay.var('_'), ret, expr) return ret