def after(): a = relay.var("a", shape=(10, 10)) b = relay.var("b", shape=(10, 10)) # add_relu function in_1 = relay.var("in_1", shape=(10, 10)) in_2 = relay.var("in_2", shape=(10, 10)) add_node = relay.add(in_1, in_2) relu_node = relay.nn.relu(add_node) add_relu = relay.Function([in_1, in_2], relu_node) add_relu = add_relu.with_attr("Composite", "test.add_relu") # merged function cb_1 = relay.annotation.compiler_begin(a, "test") cb_2 = relay.annotation.compiler_begin(b, "test") r = relay.Call(add_relu, [cb_1, cb_2]) ce_1 = relay.annotation.compiler_end(r, "test") f = relay.Function([a, b], ce_1) mod = tvm.IRModule.from_expr(f) return mod
def expected(): a = relay.var('a', shape=(10, 10)) b = relay.var('b', shape=(10, 10)) # add_relu_add function in_1 = relay.var('in_1', shape=(10, 10)) in_2 = relay.var('in_2', shape=(10, 10)) add_node = relay.add(in_1, in_2) add_node_1 = relay.add(in_1, add_node) add_node_2 = relay.add(add_node_1, add_node) add_add_add = relay.Function([in_1, in_2], add_node_2) add_add_add = add_add_add.set_attribute("Primitive", tir.IntImm("int32", 1)) add_add_add = add_add_add.set_attribute("Composite", tir.StringImm("add_add_add")) # merged function sub_node = relay.subtract(a, b) call = relay.Call(add_add_add, [sub_node, b]) return relay.Function([a, b], call)
def after_B(): inputs = [ relay.var('input_' + str(i), shape=(10, 10)) for i in range(8) ] add_relu_calls = [] for i in range(4): x = relay.var('x' + str(i)) y = relay.var('x' + str(i)) add_relu = relay.add(x, y) add_relu = relay.nn.relu(add_relu) add_relu = relay.Function([x, y], add_relu) add_relu = add_relu.with_attr('Composite', 'add_relu') add_relu_call = relay.Call(add_relu, [inputs[i * 2], inputs[i * 2 + 1]]) add_relu_calls.append(add_relu_call) add = relay.add(add_relu_calls[0], add_relu_calls[1]) sub = relay.subtract(add_relu_calls[2], add_relu_calls[3]) out = relay.multiply(add, sub) return relay.Function(inputs, out)
def expected_same_output_region(): mod = tvm.IRModule() x = relay.var("x", shape=(8, 8)) y = relay.var("y", shape=(8, 8)) z = relay.var("z", shape=(8, 8)) x0 = relay.var("x0", shape=(8, 8)) y0 = relay.var("y0", shape=(8, 8)) log = relay.log(x0) sub = x0 - y0 mul = log * sub # The partitioned graph contains log, subtract, and multiply func = relay.Function([x0, y0], mul) func = set_func_attr(func, "ccompiler", "ccompiler_0") glb_0 = relay.GlobalVar("ccompiler_0") mod[glb_0] = func add = x + y call = relay.Call(glb_0, [add, z]) main = relay.Function([x, y, z], call) mod["main"] = main return mod
def __call__(self, *args): """Call the global variable. Parameters ---------- args: List[RelayExpr] The arguments to the call. Returns ------- call: BaseExpr A call taking the variable as a function. """ # pylint: disable=import-outside-toplevel if all(isinstance(x, RelayExpr) for x in args): from tvm import relay return relay.Call(self, args) arg_types = [type(x) for x in args] raise RuntimeError( "Do not know how to handle GlobalVar.__call__ for types {}".format(arg_types))
def create_graph(): def create_external_func1(mod_, compiler_name, symbol_name): x_int = relay.var("x_int", shape=(10, 10)) z0 = relay.nn.relu(x_int) f1 = relay.Function([x_int], z0) f1 = set_func_attr(f1, compiler_name, symbol_name) glb_f1 = relay.GlobalVar(symbol_name) mod_[glb_f1] = f1 mod_ = relay.transform.InferType()(mod_) return glb_f1, mod_ mod = tvm.IRModule() x = relay.var("x", shape=(10, 10)) glb_symbol_f1, mod = create_external_func1(mod, "ethosu", "ethosu_0") r = relay.Call(glb_symbol_f1, [x]) main = relay.Function([x], r) mod["main"] = main mod = relay.transform.InferType()(mod) return mod
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_count_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(0, dtype="int32") iarg = relay.var("i", shape=[], dtype="int32") mod["main"] = relay.Function([iarg], sum_up(iarg)) for tgt, ctx in tvm.testing.enabled_targets(): result = veval(mod, i_data, ctx=ctx, target=tgt) tvm.testing.assert_allclose(result.asnumpy(), i_data) check_result([i_data], i_data, mod=mod)
def expected(): 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]) + i 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 test_runtime_module_generation(check_result): shape = (8,) x_data = np.random.randint(255, size=shape).astype("float32") y_data = np.random.randint(255, size=shape).astype("float32") inputs = {"x": x_data, "y": y_data} x0 = relay.var("x0", shape=shape, dtype="float32") y0 = relay.var("y0", shape=shape, dtype="float32") z = x0 + y0 func = relay.Function([x0, y0], z) func = set_external_func_attr(func, "example_target_hook", "replace_add_with_subtract") # Test hook to trigger TIRToRuntime code generation func = func.with_attr("tir_to_runtime", True) x = relay.var("x", shape=(8,), dtype="float32") y = relay.var("y", shape=(8,), dtype="float32") call = relay.Call(func, [x, y]) func = IRModule.from_expr(call) check_result(func, inputs, (8,), x_data * y_data)
def test_extern_dnnl(check_result): dtype = "float32" ishape = (1, 32, 14, 14) w1shape = (32, 1, 3, 3) data0 = relay.var("data0", shape=(ishape), dtype=dtype) weight0 = relay.var("weight0", shape=(w1shape), dtype=dtype) data1 = relay.var("data0", shape=(ishape), dtype=dtype) weight1 = relay.var("weight0", shape=(w1shape), dtype=dtype) weight2 = relay.var("weight1", shape=(w1shape), dtype=dtype) depthwise_conv2d_1 = relay.nn.conv2d(data1, weight1, kernel_size=(3, 3), padding=(1, 1), groups=32) depthwise_conv2d_2 = relay.nn.conv2d(depthwise_conv2d_1, weight2, kernel_size=(3, 3), padding=(1, 1), groups=32) out = relay.add(depthwise_conv2d_1, depthwise_conv2d_2) f = relay.Function([data1, weight1, weight2], out) ref_mod = tvm.IRModule() ref_mod["main"] = f f = set_external_func_attr(f, "dnnl", "dnnl_0") call = relay.Call(f, [data0, weight0, weight0]) mod = tvm.IRModule.from_expr(call) i_data = np.random.uniform(0, 1, ishape).astype(dtype) w_data = np.random.uniform(0, 1, w1shape).astype(dtype) ref_ex = relay.create_executor("graph", mod=ref_mod, device=tvm.cpu()) ref_res = ref_ex.evaluate()(i_data, w_data, w_data) check_result(mod, { "data0": i_data, "weight0": w_data }, (1, 32, 14, 14), ref_res.numpy(), tol=1e-5)
def test_extern_dnnl_const(): if not tvm.get_global_func("relay.ext.dnnl", True): print("skip because DNNL codegen is not available") return dtype = "float32" ishape = (1, 32, 14, 14) w1shape = (32, 1, 3, 3) data0 = relay.var("data0", shape=(ishape), dtype=dtype) w_data = np.random.uniform(0, 1, w1shape).astype(dtype) data1 = relay.var("data0", shape=(ishape), dtype=dtype) weight1 = relay.const(w_data, dtype=dtype) weight2 = relay.const(w_data, dtype=dtype) depthwise_conv2d_1 = relay.nn.conv2d(data1, weight1, kernel_size=(3, 3), padding=(1, 1), groups=32) depthwise_conv2d_2 = relay.nn.conv2d(depthwise_conv2d_1, weight2, kernel_size=(3, 3), padding=(1, 1), groups=32) out = relay.add(depthwise_conv2d_1, depthwise_conv2d_2) f = relay.Function([data1], out) ref_mod = tvm.IRModule() ref_mod["main"] = f f = set_external_func_attr(f, "dnnl", "dnnl_0") call = relay.Call(f, [data0]) mod = tvm.IRModule.from_expr(call) i_data = np.random.uniform(0, 1, ishape).astype(dtype) ref_ex = relay.create_executor("graph", mod=ref_mod, ctx=tvm.cpu()) ref_res = ref_ex.evaluate()(i_data) check_result(mod, {"data0": i_data}, (1, 32, 14, 14), ref_res.asnumpy(), tol=1e-5)
def create_model(): ifm_count = int(np.prod(ifm_shape)) ifm2_count = int(np.prod(ifm2_shape)) # Create a "partitioned" Relay function ifms = relay.var("ifms", shape=[ifm_count + ifm2_count], dtype=dtype) split = relay.split(ifms, [ifm_count]) ifm = relay.reshape(split[0], newshape=ifm_shape) ifm2 = relay.reshape(split[1], newshape=ifm2_shape) shr_op = infra.make_ethosu_binary_elementwise(ifm, ifm2, ifm_shape[3], ifm2_shape[3], "SHR", ofm_dtype, reversed_operands) glb_ethosu = relay.GlobalVar("tvmgen_default_ethos_u_main_0") func = (relay.Function([ifms], shr_op).with_attr( "Inline", 1).with_attr("Compiler", "ethos-u").with_attr( "global_symbol", "tvmgen_default_ethos_u_main_0").with_attr("Primitive", 1)) mod = tvm.IRModule() mod[glb_ethosu] = func mod = relay.transform.InferType()(mod) # Main ifm = relay.var("ifm", shape=ifm_shape, dtype=dtype) ifm2 = relay.var("ifm2", shape=ifm2_shape, dtype=dtype) call = relay.Call( glb_ethosu, [ relay.concatenate( data=( relay.reshape(ifm, newshape=ifm_count), relay.reshape(ifm2, newshape=ifm2_count), ), axis=0, ) ], ) mod["main"] = relay.Function([ifm, ifm2], call) mod = relay.transform.InferType()(mod) return mod
def test_cps_pe(): def destroy_ref(x): x = run_infer_type(x) x = to_cps(x) x = run_infer_type(x) y = un_cps(x) y = run_infer_type(y) x = run_opt_pass(x, transform.Sequential([transform.PartialEvaluate(), transform.DeadCodeElimination(inline_once=True)])) assert Feature.fRefCreate not in detect_feature(x) unit = relay.Function([], relay.const(0., dtype='float32')) f_ref = relay.Var("f_ref") one = relay.const(1., dtype='float32') two = relay.const(2., dtype='float32') cond = relay.var(shape=(), dtype='uint1', name_hint='cond') true_branch = relay.RefWrite(f_ref, relay.Function([], one)) false_branch = relay.RefWrite(f_ref, relay.Function([], two)) if_expr = relay.If(cond, true_branch, false_branch) stmt = relay.Let(f_ref, relay.RefCreate(unit), relay.Let(relay.Var("x"), if_expr, relay.Call(relay.RefRead(f_ref), []))) F = relay.Function([cond], stmt) destroy_ref(F) G = relay.Function([cond], relay.If(cond, one, two)) G = run_infer_type(G) G = relay.transform.gradient(G) destroy_ref(G) x = relay.var("x", shape=(1, 16)) y = relay.var("y", shape=(1, 16)) z = relay.var("z", shape=(1, 16)) cond = relay.var("cond", shape=(), dtype='uint1') H = relay.If(cond, x, y) H = relay.add(H, z) H = relay.Function([cond,x,y,z], H) H = run_infer_type(H) H = relay.transform.gradient(H) destroy_ref(H)
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["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_mixed_data_types(): """ Test a graph with a primitive function that has mixed datatypes. """ def get_inner_func(): x = relay.var("x", shape=(1, 2, 2, 2), dtype="int16") x = relay.cast(x, dtype="uint32") x = _create_primitive_function(x) return x ifm = relay.var("input", shape=(1, 2, 2, 2), dtype="int16") x = get_inner_func() x = relay.Call(x, [ifm]) mod = tvm.IRModule.from_expr(x) expected_annotations = [ [(1 * 2 * 2 * 2) * 2 + (1 * 2 * 2 * 2) * 4], ] expected_io_annotation = (1 * 2 * 2 * 2) * 2 + (1 * 2 * 2 * 2) * 4 _check_used_memory_annotations(mod, expected_annotations, expected_io_annotation)
def test_external_function(): y0_data = np.random.uniform(0, 1, (8, 8)).astype("float32") x0 = relay.var("x0", shape=(8, 8)) y0_const = relay.const(y0_data, "float32") z0 = x0 + y0_const ef = relay.Function([x0], z0, relay.TensorType((8, 8), "float32")) ev = relay.GlobalVar("external_function") ef = set_external_func_attr(ef, "cmsis-nn", ev.name_hint) x = relay.var("x", shape=(8, 8)) c = relay.Call(ev, [x]) mf = relay.Function([x], c, relay.TensorType((8, 8), "float32")) mv = relay.GlobalVar("main") mod = tvm.IRModule() mod[ev] = ef mod[mv] = mf mod = ExtractConstantsFromPartitionedFunction()(mod) CheckFunctionsForConstants().check_num_constants(mod[ev]) relay.transform.InferType()(mod)
def test_pow(): mod = relay.Module() p = Prelude(mod) shape = (10, 10) dtype = 'float32' t = relay.TensorType(shape, dtype) x = relay.var("x", t) double = relay.Function([x], x + x) i = relay.var("i", t) func = relay.Function([i], relay.Call(p.iterate(double, p.s(p.s(p.s(p.z())))), [i])) back_func = relay.ir_pass.infer_type(gradient(func, mod=mod), mod=mod) assert back_func.checked_type == relay.FuncType( [t], relay.TupleType([t, relay.TupleType([t])])) i_nd = rand(dtype, *shape) ex = create_executor(mod=mod) forward, (grad_i, ) = ex.evaluate(back_func)(i_nd) np.testing.assert_allclose(forward.asnumpy(), 8 * i_nd.asnumpy()) np.testing.assert_allclose(grad_i.asnumpy(), 8 * np.ones_like(grad_i.asnumpy()))
def expected(): mod = tvm.IRModule() x = relay.const(ones) y = relay.var("y", shape=(8, 8)) x0 = relay.const(ones) y0 = relay.var("y0", shape=(8, 8)) add = x0 + y0 # Function that uses C compiler func = relay.Function([y0], add) func = func.with_attr("Primitive", tvm.tir.IntImm("int32", 1)) func = func.with_attr("Inline", tvm.tir.IntImm("int32", 1)) func = func.with_attr("Compiler", tvm.tir.StringImm("ccompiler")) func = func.with_attr("ExternalSymbol", tvm.tir.StringImm("ccompiler_0")) glb_0 = relay.GlobalVar("ccompiler_0") mod[glb_0] = func add_call = relay.Call(glb_0, [y]) log = relay.log(add_call) main = relay.Function([y], log) mod["main"] = main return mod
def get_net(iterations, num_hidden, batch_size=1, dtype="float32"): """Constructs an unrolled RNN with LSTM cells""" input_type = relay.TensorType((batch_size, num_hidden), dtype) weight_type = relay.TensorType((4 * num_hidden, num_hidden), dtype) bias_type = relay.TensorType((4 * num_hidden,), dtype) state_type = relay.TupleType([input_type, input_type]) cell_type = relay.TupleType([input_type, state_type]) builder = relay.ScopeBuilder() zeros = builder.let(("zeros", input_type), relay.zeros((batch_size, num_hidden), dtype)) init_states = builder.let(("init_states", state_type), relay.Tuple([zeros, zeros])) states = init_states out = None for i in range(iterations): inputs = relay.Var("data", input_type) i2h_weight = relay.Var("i2h_%s_weight" % i, weight_type) i2h_bias = relay.Var("i2h_%i_bias" % i, bias_type) h2h_weight = relay.Var("h2h_%s_weight" % i, weight_type) h2h_bias = relay.Var("h2h_%s_bias" % i, bias_type) cell_fn = lstm_cell(num_hidden, batch_size, dtype, "lstm_%s" % i) call = builder.let( ("call_%s" % i, cell_type), relay.Call(cell_fn, [inputs, states, i2h_weight, i2h_bias, h2h_weight, h2h_bias]), ) new_out = builder.let(("out_%s" % i, input_type), relay.TupleGetItem(call, 0)) new_states = builder.let(("states_%s" % i, state_type), relay.TupleGetItem(call, 1)) states = new_states out = new_out builder.ret(out) body = builder.get() args = relay.analysis.free_vars(body) return relay.Function(args, body, input_type)
def test_extern_gcc_consts(check_result): shape = (8, 8) dtype = "float32" x = relay.var("x", shape=shape) y0_data = np.random.uniform(0, 1, shape).astype(dtype) x0 = relay.var("x0", shape=shape) y0_const = relay.const(y0_data, dtype) z = x0 + y0_const f = relay.Function([x0], z) f = set_external_func_attr(f, "ccompiler", "ccompiler_0") call = relay.Call(f, [x]) mod = tvm.IRModule.from_expr(call) # Note that while the VMCompiler get_params() will return all 'parameters' from both # TVM and external codegen compiled code, the GraphExecutor.get_params() will return only # those from non-external modules. So in the following we'll test by execution rather than # test by inspection. x_data = np.random.rand(*shape).astype(dtype) inputs = {"x": x_data} expected_result = x_data + y0_data check_result(mod, inputs, shape, expected_result, target="llvm")
def expected(): x = relay.var('x', shape=(1, 8)) beta = relay.var("beta", shape=(8, )) gamma = relay.var("gamma", shape=(8, )) moving_mean = relay.var("moving_mean", shape=(8, )) moving_var = relay.var("moving_var", shape=(8, )) # bn_relu function in_1 = relay.var('x1', shape=(1, 8)) in_2 = relay.var('gamma1', shape=(8, )) in_3 = relay.var('beta1', shape=(8, )) in_4 = relay.var('moving_mean1', shape=(8, )) in_5 = relay.var('moving_var1', shape=(8, )) bn_node = relay.nn.batch_norm(in_1, in_2, in_3, in_4, in_5) tuple_get_item_node = bn_node[0] relu_node = relay.nn.relu(tuple_get_item_node) bn_relu = relay.Function([in_1, in_2, in_3, in_4, in_5], relu_node) bn_relu = bn_relu.with_attr("Composite", tir.StringImm("bn_relu")) # merged function r = relay.Call(bn_relu, [x, gamma, beta, moving_mean, moving_var]) return relay.Function([x, gamma, beta, moving_mean, moving_var], r)
def expected(): def create_external_func1(mod_, compiler_name, symbol_name): x_int = relay.var("x_int", shape=(10, 10)) p0 = relay.nn.relu(x_int) q0 = relay.tanh(x_int) # reshapes p0_reshaped = relay.reshape(p0, newshape=100) q0_reshaped = relay.reshape(q0, newshape=100) ofms = relay.concatenate((p0_reshaped, q0_reshaped), 0) f1 = relay.Function([x_int], ofms) f1 = set_func_attr(f1, compiler_name, symbol_name) glb_f1 = relay.GlobalVar(symbol_name) mod_[glb_f1] = f1 mod_ = relay.transform.InferType()(mod_) return glb_f1, mod_ mod = tvm.IRModule() x = relay.var("x", shape=(10, 10)) glb_symbol_f1, mod = create_external_func1(mod, "ethosu", "ethosu_0") ofms = relay.Call(glb_symbol_f1, [x]) # splits (p0_flat, q0_flat) = relay.split(ofms, [100]) # reshapes p0_flat_reshaped = relay.reshape(p0_flat, newshape=(10, 10)) q0_flat_reshaped = relay.reshape(q0_flat, newshape=(10, 10)) # original output tuple_out = relay.Tuple([p0_flat_reshaped, q0_flat_reshaped]) p0 = relay.TupleGetItem(tuple_out, 0) q0 = relay.TupleGetItem(tuple_out, 1) r = relay.concatenate((p0, q0), axis=0) main = relay.Function([x], r) mod["main"] = main mod = relay.transform.InferType()(mod) return mod
def _full_offload(): mod = tvm.IRModule({}) # NPU function x = relay.var("x", shape=(1, 4, 4, 16), dtype="int8") max_pool = relay.nn.max_pool2d(x) composite_func = relay.Function([x], max_pool) composite_func = composite_func.with_attr("Composite", "ethos-u.pooling") inp = relay.var("input", shape=(1, 4, 4, 16), dtype="int8") compiler_func = relay.Function([inp], composite_func) compiler_func = compiler_func.with_attr("used_memory", [256 + 256]) npu_compiler_func = compiler_func.with_attr("Compiler", "ethos-u") g1 = relay.GlobalVar("g1") mod[g1] = npu_compiler_func # Main inp = relay.var("main_input", shape=(1, 4, 4, 16), dtype="int8") call = relay.Call(g1, [inp]) main_func = relay.Function([inp], call) main_func = main_func.with_attr("io_used_memory", 256 + 256) mod["main"] = main_func return 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)))
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 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([i_data, accum_data], sum(range(1, loop_bound + 1)), mod=mod)
def after_B(): inputs = [ relay.var("input_" + str(i), shape=(10, 10)) for i in range(8) ] add_relu_calls = [] for i in range(4): x = relay.var("x" + str(i)) y = relay.var("x" + str(i)) add_relu = relay.add(x, y) add_relu = relay.nn.relu(add_relu) add_relu = relay.Function([x, y], add_relu) add_relu = add_relu.with_attr("Composite", "add_relu") add_relu = add_relu.with_attr("PartitionedFromPattern", "add_nn.relu_") add_relu_call = relay.Call(add_relu, [inputs[i * 2], inputs[i * 2 + 1]]) add_relu_calls.append(add_relu_call) add = relay.add(add_relu_calls[0], add_relu_calls[1]) sub = relay.subtract(add_relu_calls[2], add_relu_calls[3]) out = relay.multiply(add, sub) return relay.Function(inputs, out)
def expected(dshape): p0 = relay.var("p0", shape=dshape) c = conv(p0) f0 = relay.Function(relay.analysis.free_vars(c), c) f0 = f0.with_attr("Primitive", tvm.tir.IntImm("int32", 1)) p01 = relay.var("p01", shape=dshape) c = conv(p01) f1 = relay.Function(relay.analysis.free_vars(c), c) f1 = f1.with_attr("Primitive", tvm.tir.IntImm("int32", 1)) p02 = relay.var("p02", shape=dshape) p12 = relay.var("p12", shape=dshape) concat1 = relay.concatenate((p02, p12), axis=1) f_concat1 = relay.Function([p02, p12], concat1) f_concat1 = f_concat1.with_attr("Primitive", tvm.tir.IntImm("int32", 1)) dshape2 = (dshape[0], dshape[1] * 2, dshape[2], dshape[3]) p03 = relay.var("p03", shape=dshape2) c = conv(p03) f2 = relay.Function(relay.analysis.free_vars(c), c) f2 = f2.with_attr("Primitive", tvm.tir.IntImm("int32", 1)) p04 = relay.var("p04", shape=dshape2) c = conv(p04) f3 = relay.Function(relay.analysis.free_vars(c), c) f3 = f3.with_attr("Primitive", tvm.tir.IntImm("int32", 1)) p05 = relay.var("p05", shape=dshape) p15 = relay.var("p15", shape=dshape) concat2 = relay.concatenate((p05, p15), axis=1) f_concat2 = relay.Function([p05, p15], concat2) f_concat2 = f_concat2.with_attr("Primitive", tvm.tir.IntImm("int32", 1)) x = relay.var("x", shape=dshape) c1 = relay.Call(f0, [x, relay.var("w1")]) c2 = relay.Call(f1, [x, relay.var("w2")]) concat = relay.Call(f_concat1, [c1, c2]) c3 = relay.Call(f2, [concat, relay.var("w3")]) c4 = relay.Call(f3, [concat, relay.var("w4")]) out = relay.Call(f_concat2, [c3, c4]) return relay.Function(relay.analysis.free_vars(out), out)
def visit_call(self, call): # TODO(weberlo) use graph pattern matching? if not hasattr(call.op, "name") or call.op.name not in ALLOWED_CONVERSION_OPS: new_args = [] for arg in call.args: new_arg = self.visit(arg) if len(self.subtree_params) == 0: new_args.append(new_arg) else: assert len(self.subtree_params) == 1 param = next(iter(self.subtree_params)) pre_param = self.prefix_sb.let(param.name_hint, new_arg) self.subtree_params.clear() mid_param = relay.Var(param.name_hint, arg.checked_type) self.prefix_binding_map[mid_param] = pre_param # return new parameter, then we can use # relay.analysis.free_vars at the end of the pass to generate # new `mid_func` type signature new_args.append(mid_param) return relay.Call(call.op, new_args, call.attrs) return super().visit_call(call)
def make_partitioned_function(relay_op): ifm0 = relay.analysis.free_vars(relay_op) ifm_shape = ifm0[0].type_annotation.shape ifm = relay.var("ifm", shape=ifm_shape, dtype="int8") glb_ethosu = relay.GlobalVar("tvmgen_default_ethosu_main_0") func = (relay.Function(ifm0, relay_op).with_attr("Inline", 1).with_attr( "Compiler", "ethos-u").with_attr("global_symbol", "tvmgen_default_ethosu_main_0").with_attr( "Primitive", 1)) mod = tvm.IRModule() mod[glb_ethosu] = func mod = relay.transform.InferType()(mod) call = relay.Call(glb_ethosu, [ifm]) mod["main"] = relay.Function([ifm], call) mod = relay.transform.InferType()(mod) return mod