def test_recursive_func(): mod = tvm.IRModule({}) x = relay.var('x', shape=[], dtype='int32') fn0 = relay.Function([x], x) gx = relay.GlobalVar("gx") mod[gx] = fn0 sum_up = relay.GlobalVar('sum_up') i = relay.var('i', shape=[], dtype='int32') sb = relay.ScopeBuilder() with sb.if_scope(relay.equal(i, relay.const(0, dtype='int32'))): sb.ret(i) with sb.else_scope(): one_less = relay.subtract(i, relay.const(1, dtype='int32')) global_call = gx(i) rec_call = relay.Call(sum_up, [one_less]) + global_call sb.ret(relay.add(rec_call, i)) func = relay.Function([i], sb.get(), ret_type=relay.TensorType([], 'int32')) func = func.set_attribute("Compiler", tvm.tir.StringImm("a")) mod[sum_up] = func iarg = relay.var('i', shape=[], dtype='int32') mod["main"] = relay.Function([iarg], sum_up(iarg)) call_graph = relay.CallGraph(mod) assert call_graph.is_recursive(sum_up) assert call_graph.ref_count(sum_up) == 2 assert call_graph.ref_count(gx) == 1 assert call_graph.ref_count("main") == 0
def get_mod(): mod = tvm.IRModule({}) x = relay.var('x', shape=[], dtype='int32') fn0 = relay.Function([x], x) fn0 = fn0.set_attribute("Inline", tvm.tir.IntImm("int32", 1)) gx = relay.GlobalVar("gx") mod[gx] = fn0 sum_up = relay.GlobalVar('sum_up') i = relay.var('i', shape=[], dtype='int32') sb = relay.ScopeBuilder() with sb.if_scope(relay.equal(i, relay.const(0, dtype="int32"))): sb.ret(i) with sb.else_scope(): one_less = relay.subtract(i, relay.const(1, dtype="int32")) global_call = gx(i) rec_call = relay.Call(sum_up, [one_less]) + global_call sb.ret(relay.add(rec_call, i)) func = relay.Function([i], sb.get(), ret_type=relay.TensorType([], "int32")) func = func.set_attribute("Inline", tvm.tir.IntImm("int32", 1)) mod[sum_up] = func iarg = relay.var("i", shape=[], dtype='int32') mod["main"] = relay.Function([iarg], sum_up(iarg)) return mod
def after(): data = relay.var("data", shape=(1, 32)) eq1 = relay.var("e1", shape=[], dtype="float32") eq2 = relay.var("e2", shape=[], dtype="float32") cb_1 = relay.annotation.compiler_begin(eq1, target) cb_2 = relay.annotation.compiler_begin(eq2, target) equality_condition = relay.equal(cb_1, cb_2) ce_1 = relay.annotation.compiler_end(equality_condition, target) # if condition cb_3 = relay.annotation.compiler_begin(data, target) true_branch = relay.tanh(cb_3) ce_2 = relay.annotation.compiler_end(true_branch, target) # else condition cb_4 = relay.annotation.compiler_begin(data, target) false_branch = relay.sigmoid(cb_4) ce_3 = relay.annotation.compiler_end(false_branch, target) if_condition = relay.If(ce_1, ce_2, ce_3) cb_5 = relay.annotation.compiler_begin(if_condition, target) erf_out = relay.erf(cb_5) ce_4 = relay.annotation.compiler_end(erf_out, target) func = relay.Function([data, eq1, eq2], ce_4) mod = tvm.IRModule.from_expr(func) return mod
def test_equal(): i = relay.var("i", shape=[], dtype="int32") j = relay.var("i", shape=[], dtype="int32") z = relay.equal(i, j) func = relay.Function([i, j], z, ret_type=relay.TensorType([], "bool")) i_data = relay.const(0, "int32") j_data = relay.const(0, "int32")
def _make_env_find(self, m, rval_t): ctr = m['ctr'] gv = relay.GlobalVar(f"$_env_find<{ctr.name_hint}>") env = relay.Var("env", env_type(env_val())) key = relay.Var("key", relay.ty.scalar_type('int64')) dft = relay.Var("dft", rval_t) k = relay.Var("k") v = relay.Var("v") r = relay.Var("r") x = relay.Var("x") extract_clause = adt.Clause( adt.PatternConstructor(ctr, [adt.PatternVar(x)]), x) empty_clause = adt.Clause(adt.PatternConstructor(empty_env, []), dft) cons_clause = adt.Clause( adt.PatternConstructor( cons_env, [adt.PatternVar(k), adt.PatternVar(v), adt.PatternVar(r)]), relay.If(relay.equal(key, k), adt.Match(v, [extract_clause], complete=False), relay.Call(gv, [r, key, dft]))) body = adt.Match(env, [empty_clause, cons_clause]) fn = relay.Function([env, key, dft], body, rval_t) m['env_find'] = (gv, fn) return gv, fn
def test_recursion(): """ Program: def @f(%n: int32, %data: float32) -> float32 { if (%n == 0) { %data } else { @f(%n - 1, log(%data)) } } """ sb = relay.ScopeBuilder() f = relay.GlobalVar("f") ti32 = relay.scalar_type("int32") tf32 = relay.scalar_type("float32") n = relay.var("n", ti32) data = relay.var("data", tf32) with sb.if_scope(relay.equal(n, relay.const(0, ti32))): sb.ret(data) with sb.else_scope(): sb.ret(f(relay.subtract(n, relay.const(1, ti32)), relay.log(data))) mod = tvm.IRModule() mod[f] = relay.Function([n, data], sb.get()) mod = infer_mod(mod) assert "@f(%1, %2)" in mod.astext() assert mod["f"].checked_type == relay.FuncType([ti32, tf32], tf32)
def test_loop(): mod = tvm.IRModule({}) sum_up = relay.GlobalVar("sum_up") i = relay.var("i", shape=[], dtype="int32") accum = relay.var("accum", shape=[], dtype="int32") sb = ScopeBuilder() with sb.if_scope(relay.equal(i, relay.const(0, "int32"))): sb.ret(accum) with sb.else_scope(): one_less = relay.subtract(i, relay.const(1, "int32")) new_accum = relay.add(accum, i) sb.ret(relay.Call(sum_up, [one_less, new_accum])) func = relay.Function([i, accum], sb.get()) mod[sum_up] = func mod = transform.InferType()(mod) loop_bound = 0 i_data = np.array(loop_bound, dtype="int32") accum_data = np.array(0, dtype="int32") iarg = relay.var("i", shape=[], dtype="int32") aarg = relay.var("accum", shape=[], dtype="int32") mod["main"] = relay.Function([iarg, aarg], sum_up(iarg, aarg)) result = get_serialized_output(mod, i_data, accum_data) tvm.testing.assert_allclose(result.asnumpy(), sum(range(1, loop_bound + 1)))
def test_recursion(): """ Program: def @f(%n: int32, %data: float32) -> float32 { if (%n == 0) { %data } else { @f(%n - 1, log(%data)) } } """ sb = relay.ScopeBuilder() f = relay.GlobalVar("f") ti32 = relay.scalar_type("int32") tf32 = relay.scalar_type("float32") n = relay.var("n", ti32) data = relay.var("data", tf32) with sb.if_scope(relay.equal(n, relay.const(0, ti32))): sb.ret(data) with sb.else_scope(): sb.ret(f(relay.subtract(n, relay.const(1, ti32)), relay.log(data))) mod = relay.Module() mod[f] = relay.Function([n, data], sb.get()) assert "@f(%1, %2) /* ty=float32 */" in mod.astext() assert mod[f].checked_type == relay.FuncType([ti32, tf32], tf32)
def _make_env_update(self, m, rval_t): ctr = m['ctr'] gv = relay.GlobalVar(f"$_env_update<{ctr.name_hint}>") env = relay.Var("env", env_type(env_val())) key = relay.Var("key", relay.ty.scalar_type('int64')) val = relay.Var("val", rval_t) k = relay.Var("k") v = relay.Var("v") r = relay.Var("r") empty_clause = adt.Clause(adt.PatternConstructor(empty_env, []), cons_env(key, ctr(val), env)) cons_clause = adt.Clause( adt.PatternConstructor( cons_env, [adt.PatternVar(k), adt.PatternVar(v), adt.PatternVar(r)]), relay.If(relay.equal(key, k), cons_env(key, ctr(val), env), cons_env(k, v, relay.Call(gv, [r, key, val])))) body = adt.Match(env, [empty_clause, cons_clause]) fn = relay.Function([env, key, val], body, env_type(env_val())) m['env_update'] = (gv, fn) return gv, fn
def test_recursion(): """ Program: let sum_twice(n: i32) -> i32 = { m = (n * 2) if (n == 0) { return m; } else { return m + sum(n - 1); } } sum_twice(5); """ return # cannot be run as fuse_ops need to recursively visit mod = relay.Module() i64 = relay.TensorType((), 'int64') f = relay.GlobalVar("f") n = relay.Var("n", i64) m = n * relay.const(2, 'int64') funcbody = relay.If(relay.equal(n, relay.const(0, 'int64')), m, m + f(n - relay.const(1, 'int64'))) value = relay.Function([n], funcbody, i64, []) mod[f] = value check_eval(f(relay.const(5, 'int64')), 30.0, mod=mod) old_f = mod[f] mod = transform.ToANormalForm()(mod) f = mod[f] check_eval(f(relay.const(5, 'int64')), 30.0, mod=mod)
def test_loop(): mod = relay.module.Module({}) sum_up = relay.GlobalVar('sum_up') i = relay.var('i', shape=[], dtype='int32') accum = relay.var('accum', shape=[], dtype='int32') sb = ScopeBuilder() with sb.if_scope(relay.equal(i, relay.const(0, 'int32'))): sb.ret(accum) with sb.else_scope(): one_less = relay.subtract(i, relay.const(1, 'int32')) new_accum = relay.add(accum, i) sb.ret(relay.Call(sum_up, [one_less, new_accum])) func = relay.Function([i, accum], sb.get()) mod[sum_up] = func loop_bound = 0 i_data = np.array(loop_bound, dtype='int32') accum_data = np.array(0, dtype='int32') iarg = relay.var('i', shape=[], dtype='int32') aarg = relay.var('accum', shape=[], dtype='int32') mod["main"] = relay.Function([iarg, aarg], sum_up(iarg, aarg)) vm = create_vm(mod) ser = serializer.Serializer(vm) code, lib = ser.serialize() deser = deserializer.Deserializer(code, lib) des_vm = deser.deserialize() result = veval(des_vm, i_data, accum_data) tvm.testing.assert_allclose(result.asnumpy(), sum(range(1, loop_bound + 1)))
def test_recursion(): """ Program: let f(n: i32) -> i32 = { m = (n * 2) if (n == 0) { return m; } else { return m + f(n - 1); } } f(5); """ mod = tvm.IRModule() i64 = relay.TensorType((), 'int64') f = relay.GlobalVar("f") n = relay.Var("n", i64) m = n * relay.const(2, 'int64') funcbody = relay.If(relay.equal(n, relay.const(0, 'int64')), m, m + f(n - relay.const(1, 'int64'))) value = relay.Function([n], funcbody, i64, []) mod[f] = value check_eval(f(relay.const(5, 'int64')), 30.0, mod=mod) old_f = mod[f] mod = transform.ToANormalForm()(mod) f = mod[f] check_eval(f(relay.const(5, 'int64')), 30.0, mod=mod)
def test_loop(): mod = relay.module.Module({}) sum_up = relay.GlobalVar('sum_up') i = relay.var('i', shape=[], dtype='int32') accum = relay.var('accum', shape=[], dtype='int32') sb = ScopeBuilder() with sb.if_scope(relay.equal(i, relay.const(0, 'int32'))): sb.ret(accum) with sb.else_scope(): one_less = relay.subtract(i, relay.const(1, 'int32')) new_accum = relay.add(accum, i) sb.ret(relay.Call(sum_up, [one_less, new_accum])) func = relay.Function([i, accum], sb.get()) mod[sum_up] = func loop_bound = 0 i_data = np.array(loop_bound, dtype='int32') accum_data = np.array(0, dtype='int32') iarg = relay.var('i', shape=[], dtype='int32') aarg = relay.var('accum', shape=[], dtype='int32') mod["main"] = relay.Function([iarg, aarg], sum_up(iarg, aarg)) exe = create_exec(mod) code, lib = exe.save() des_exec = _vm.Executable.load_exec(code, lib) des_vm = _vm.VirtualMachine(des_exec) des_vm.init(tvm.cpu()) result = veval(des_vm, i_data, accum_data) tvm.testing.assert_allclose(result.asnumpy(), sum(range(1, loop_bound + 1)))
def test_recursion(): """ Program: let f(n: i32) -> i32 = { m = (n * 2) if (n == 0) { return m; } else { return m + f(n - 1); } } f(5); """ mod = tvm.IRModule() i64 = relay.TensorType((), "int64") f = relay.GlobalVar("f") n = relay.Var("n", i64) m = n * relay.const(2, "int64") cond = relay.equal(n, relay.const(0, "int64")) false_branch = m + f(n - relay.const(1, "int64")) funcbody = relay.If(cond, m, false_branch) value = relay.Function([n], funcbody, i64, []) mod[f] = value check_eval(f(relay.const(5, "int64")), 30.0, mod=mod) old_f = mod[f] mod = transform.ToBasicBlockNormalForm()(mod) f = mod[f] check_eval(f(relay.const(5, "int64")), 30.0, mod=mod) check_basic_block_normal_form(f)
def test_sum_loop(target, dev): mod = tvm.IRModule({}) sum_up = relay.GlobalVar("sum_up") i = relay.var("i", shape=[], dtype="int32") accum = relay.var("accum", shape=[], dtype="int32") sb = ScopeBuilder() with sb.if_scope(relay.equal(i, relay.const(0, "int32"))): sb.ret(accum) with sb.else_scope(): one_less = relay.subtract(i, relay.const(1, "int32")) new_accum = relay.add(accum, i) sb.ret(relay.Call(sum_up, [one_less, new_accum])) func = relay.Function([i, accum], sb.get()) mod[sum_up] = func mod = relay.transform.InferType()(mod) loop_bound = 0 i_data = np.array(loop_bound, dtype="int32") accum_data = np.array(0, dtype="int32") iarg = relay.var("i", shape=[], dtype="int32") aarg = relay.var("accum", shape=[], dtype="int32") mod["main"] = relay.Function([iarg, aarg], sum_up(iarg, aarg)) check_result(target, dev, [i_data, accum_data], sum(range(1, loop_bound + 1)), mod=mod)
def test_recursion(): """ Program: def f(n: i32, data: f32) -> f32 { if (n == 0) { return data; } else { return f(n - 1, log(data)); } } """ sb = relay.ScopeBuilder() f = relay.GlobalVar("f") ti32 = relay.scalar_type("int32") tf32 = relay.scalar_type("float32") n = relay.var("n", ti32) data = relay.var("data", tf32) with sb.if_scope(relay.equal(n, relay.const(0, ti32))): sb.ret(data) with sb.else_scope(): sb.ret(f(relay.subtract(n, relay.const(1, ti32)), relay.log(data))) mod = relay.Module() mod[f] = relay.Function([n, data], sb.get()) assert "%3 = @f(%1, %2)" in mod.astext() assert mod[f].checked_type == relay.FuncType([ti32, tf32], tf32)
def test_equal(): i = relay.var('i', shape=[], dtype='int32') j = relay.var('i', shape=[], dtype='int32') z = relay.equal(i, j) func = relay.Function([i, j], z, ret_type=relay.TensorType([], 'bool')) i_data = relay.const(0) j_data = relay.const(0) check_eval(func, [i_data, j_data], True)
def test_equal(): i = relay.var('i', shape=[], dtype='int32') j = relay.var('i', shape=[], dtype='int32') z = relay.equal(i, j) func = relay.Function([i, j], z, ret_type=relay.TensorType([], 'bool')) i_data = relay.const(0, 'int32') j_data = relay.const(0, 'int32') check_eval(func, [i_data, j_data], True)
def if_expr(x): '\n free_var %x: float32\n %0 = equal(%x, 2f);\n if (%0) {\n %1 = add(%x, 1f);\n multiply(%1, 2f)\n } else {\n multiply(%1, 1f)\n }\n ' one = relay.const(1, dtype='float32') two = relay.const(2, dtype='float32') v1 = relay.add(x, one) 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) return body
def expected_if_expr(x): '\n free_var %x: float32\n let %v1: float32 = add(%x, 1f /* ty=float32 */) /* ty=float32 */;\n %0 = equal(%x, 2f /* ty=float32 */) /* ty=bool */;\n if (%0) {\n multiply(%v1, 2f /* ty=float32 */) /* ty=float32 */\n } else {\n multiply(%v1, 1f /* ty=float32 */) /* ty=float32 */\n }\n ' 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 get_func_with_control_flow(): data = relay.var("data", shape=(1, 3, 224, 224)) weight = relay.var("weight", shape=(32, 3, 3, 3)) eq1 = relay.var("e1", shape=[], dtype="float32") eq2 = relay.var("e2", shape=[], dtype="float32") eq = relay.equal(eq1, eq2) true_branch = relay.zeros(shape=(1, 32, 222, 222), dtype="float32") false_branch = relay.nn.conv2d(data, weight, kernel_size=(3, 3), channels=32) ife = relay.If(eq, true_branch, false_branch) out = relay.erf(ife) return relay.Function([data, weight, eq1, eq2], out)
def before(): data = relay.var("data", shape=(1, 32)) eq1 = relay.var("e1", shape=[], dtype="float32") eq2 = relay.var("e2", shape=[], dtype="float32") eq = relay.equal(eq1, eq2) true_branch = relay.tanh(data) false_branch = relay.sigmoid(data) ife = relay.If(eq, true_branch, false_branch) out = relay.erf(ife) func = relay.Function([data, eq1, eq2], out) mod = tvm.IRModule.from_expr(func) return mod
def test_simple_loop(): mod = tvm.IRModule({}) sum_up = relay.GlobalVar("sum_up") i = relay.var("i", shape=[], dtype="int32") sb = ScopeBuilder() with sb.if_scope(relay.equal(i, relay.const(0, dtype="int32"))): sb.ret(i) with sb.else_scope(): one_less = relay.subtract(i, relay.const(1, dtype="int32")) rec_call = relay.Call(sum_up, [one_less]) sb.ret(relay.add(rec_call, i)) func = relay.Function([i], sb.get(), ret_type=relay.TensorType([], "int32")) mod[sum_up] = func i_data = np.array(10, dtype="int32") check_eval(sum_up, [i_data], sum(range(1, 11)), mod=mod)
def test_simple_loop(): mod = relay.module.Module({}) sum_up = relay.GlobalVar('sum_up') i = relay.var('i', shape=[], dtype='int32') sb = ScopeBuilder() with sb.if_scope(relay.equal(i, relay.const(0, dtype='int32'))): sb.ret(i) with sb.else_scope(): one_less = relay.subtract(i, relay.const(1, dtype='int32')) rec_call = relay.Call(sum_up, [one_less]) sb.ret(relay.add(rec_call, i)) func = relay.Function([i], sb.get(), ret_type=relay.TensorType([], 'int32')) mod[sum_up] = func i_data = np.array(10, dtype='int32') check_eval(sum_up, [i_data], sum(range(1, 11)), mod=mod)
def test_loop(): mod = relay.module.Module({}) sum_up = relay.GlobalVar('sum_up') i = relay.var('i', shape=[], dtype='int32') accum = relay.var('accum', shape=[], dtype='int32') sb = ScopeBuilder() with sb.if_scope(relay.equal(i, relay.const(0, 'int32'))): sb.ret(accum) with sb.else_scope(): one_less = relay.subtract(i, relay.const(1, 'int32')) new_accum = relay.add(accum, i) sb.ret(relay.Call(sum_up, [one_less, new_accum])) func = relay.Function([i, accum], sb.get()) mod[sum_up] = func i_data = np.array(10, dtype='int32') accum_data = np.array(0, dtype='int32') check_eval(sum_up, [i_data, accum_data], sum(range(1, 11)), mod=mod)
def get_recursive_count_loop(): mod = tvm.IRModule({}) sum_up = relay.GlobalVar("sum_up") i = relay.var("i", shape=[], dtype="int32") sb = relay.ScopeBuilder() with sb.if_scope(relay.equal(i, relay.const(0, dtype="int32"))): sb.ret(i) with sb.else_scope(): one_less = relay.subtract(i, relay.const(1, dtype="int32")) rec_call = relay.Call(sum_up, [one_less]) sb.ret(relay.add(rec_call, i)) func = relay.Function([i], sb.get(), ret_type=relay.TensorType([], "int32")) func = func.with_attr("Inline", tvm.tir.IntImm("int32", 1)) mod[sum_up] = func iarg = relay.var("i", shape=[], dtype="int32") mod["main"] = relay.Function([iarg], sum_up(iarg)) return mod, sum_up
def test_loop(): mod = relay.module.Module({}) sum_up = relay.GlobalVar('sum_up') i = relay.var('i', shape=[], dtype='int32') accum = relay.var('accum', shape=[], dtype='int32') sb = ScopeBuilder() with sb.if_scope(relay.equal(i, relay.const(0))): sb.ret(accum) with sb.else_scope(): one_less = relay.subtract(i, relay.const(1)) new_accum = relay.add(accum, i) sb.ret(relay.Call(sum_up, [one_less, new_accum])) func = relay.Function([i, accum], sb.get()) mod[sum_up] = func i_data = np.array(10, dtype='int32') accum_data = np.array(0, dtype='int32') check_eval(sum_up, [i_data, accum_data], sum(range(1, 11)), mod=mod)
def relay_take_grad_inp(c, _nb_indices, _indices, _values): assert _nb_indices.is_constant(int) values = c.ref(_values) r_indices = relay.reshape(c.ref(_indices), tuple(_indices.abstract.xshape()) + (1, )) n_rows = _nb_indices.value n_cols = _values.abstract.xshape()[-1] outputs = [] indices_dtype = type_to_np_dtype(_indices.abstract.element.xtype()) out_dtype = type_to_np_dtype(_values.abstract.element.xtype()) for i in range(n_rows): select_entries = relay.equal(r_indices, relay.const(i, indices_dtype)) casted_select = relay.cast(select_entries, out_dtype) select_dout = relay.multiply(casted_select, values) reshape_out = relay.reshape(select_dout, (-1, n_cols)) vector = relay.sum(reshape_out, 0) outputs.append(relay.reshape(vector, (1, n_cols))) return relay.concatenate(outputs, 0)
def test_count_loop(): mod = relay.module.Module({}) sum_up = relay.GlobalVar('sum_up') i = relay.var('i', shape=[], dtype='int32') sb = ScopeBuilder() with sb.if_scope(relay.equal(i, relay.const(0, dtype='int32'))): sb.ret(i) with sb.else_scope(): one_less = relay.subtract(i, relay.const(1, dtype='int32')) rec_call = relay.Call(sum_up, [one_less]) sb.ret(relay.add(rec_call, i)) func = relay.Function([i], sb.get(), ret_type=relay.TensorType([], 'int32')) mod[sum_up] = func i_data = np.array(0, dtype='int32') iarg = relay.var('i', shape=[], dtype='int32') mod[mod.entry_func] = relay.Function([iarg], sum_up(iarg)) result = veval(mod, i_data) tvm.testing.assert_allclose(result.asnumpy(), i_data)
def test_recursion(): '\n Program:\n let f(n: i32) -> i32 = {\n m = (n * 2)\n if (n == 0) {\n return m;\n } else {\n return m + f(n - 1);\n }\n }\n f(5);\n ' mod = tvm.IRModule() i64 = relay.TensorType((), 'int64') f = relay.GlobalVar('f') n = relay.Var('n', i64) m = (n * relay.const(2, 'int64')) cond = relay.equal(n, relay.const(0, 'int64')) false_branch = (m + f((n - relay.const(1, 'int64')))) funcbody = relay.If(cond, m, false_branch) value = relay.Function([n], funcbody, i64, []) mod[f] = value check_eval(f(relay.const(5, 'int64')), 30.0, mod=mod) old_f = mod[f] mod = transform.ToBasicBlockNormalForm()(mod) f = mod[f] check_eval(f(relay.const(5, 'int64')), 30.0, mod=mod) check_basic_block_normal_form(f)
def if_expr(x): """ free_var %x: float32 %0 = equal(%x, 2f); if (%0) { %1 = add(%x, 1f); multiply(%1, 2f) } else { multiply(%1, 1f) } """ one = relay.const(1, dtype="float32") two = relay.const(2, dtype="float32") v1 = relay.add(x, one) 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) return body
def test_count_loop(target, dev): mod = tvm.IRModule({}) sum_up = relay.GlobalVar("sum_up") i = relay.var("i", shape=[], dtype="int32") sb = ScopeBuilder() with sb.if_scope(relay.equal(i, relay.const(0, dtype="int32"))): sb.ret(i) with sb.else_scope(): one_less = relay.subtract(i, relay.const(1, dtype="int32")) rec_call = relay.Call(sum_up, [one_less]) sb.ret(relay.add(rec_call, i)) func = relay.Function([i], sb.get(), ret_type=relay.TensorType([], "int32")) mod[sum_up] = func i_data = np.array(0, dtype="int32") iarg = relay.var("i", shape=[], dtype="int32") mod["main"] = relay.Function([iarg], sum_up(iarg)) result = veval(mod, i_data, device=dev, target=target) tvm.testing.assert_allclose(result.numpy(), i_data) check_result(target, dev, [i_data], i_data, 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_sum_loop(): mod = tvm.IRModule({}) sum_up = relay.GlobalVar('sum_up') i = relay.var('i', shape=[], dtype='int32') accum = relay.var('accum', shape=[], dtype='int32') sb = ScopeBuilder() with sb.if_scope(relay.equal(i, relay.const(0, 'int32'))): sb.ret(accum) with sb.else_scope(): one_less = relay.subtract(i, relay.const(1, 'int32')) new_accum = relay.add(accum, i) sb.ret(relay.Call(sum_up, [one_less, new_accum])) func = relay.Function([i, accum], sb.get()) mod[sum_up] = func loop_bound = 0 i_data = np.array(loop_bound, dtype='int32') accum_data = np.array(0, dtype='int32') iarg = relay.var('i', shape=[], dtype='int32') aarg = relay.var('accum', shape=[], dtype='int32') mod["main"] = relay.Function([iarg, aarg], sum_up(iarg, aarg)) check_result([i_data, accum_data], sum(range(1, loop_bound + 1)), mod=mod)
def test_sum_loop(): mod = relay.module.Module({}) sum_up = relay.GlobalVar('sum_up') i = relay.var('i', shape=[], dtype='int32') accum = relay.var('accum', shape=[], dtype='int32') sb = ScopeBuilder() with sb.if_scope(relay.equal(i, relay.const(0, 'int32'))): sb.ret(accum) with sb.else_scope(): one_less = relay.subtract(i, relay.const(1, 'int32')) new_accum = relay.add(accum, i) sb.ret(relay.Call(sum_up, [one_less, new_accum])) func = relay.Function([i, accum], sb.get()) mod[sum_up] = func loop_bound = 0 i_data = np.array(loop_bound, dtype='int32') accum_data = np.array(0, dtype='int32') iarg = relay.var('i', shape=[], dtype='int32') aarg = relay.var('accum', shape=[], dtype='int32') mod[mod.entry_func] = relay.Function([iarg, aarg], sum_up(iarg, aarg)) result = veval(mod, i_data, accum_data) tvm.testing.assert_allclose(result.asnumpy(), sum(range(1, loop_bound + 1)))