def partition_for_ethosu(
mod: tvm.ir.IRModule, params: Optional[Dict[str, tvm.runtime.NDArray]] = None, **opts
):
"""This helper function partition the relay graph as produced by the
relay frontend for a given model into external functions
to be presented to the codegen.
Parameters
----------
mod : tvm.ir.IRModule
The IRModule that gets generated from a relay frontend
params : Optional[Dict[str, tvm.runtime.NDArray]]
Constant input parameters.
Returns
-------
mod : IRModule
The partitioned IRModule with external global functions
"""
if params:
mod["main"] = bind_params_by_name(mod["main"], params)
pattern = relay.op.contrib.get_pattern_table("ethos-u")
mod = relay.transform.InferType()(mod)
mod = relay.transform.MergeComposite(pattern)(mod)
mod = relay.transform.AnnotateTarget("ethos-u")(mod)
mod = relay.transform.MergeCompilerRegions()(mod)
mod = relay.transform.InferType()(mod)
mod = relay.transform.PartitionGraph()(mod)
mod = relay.transform.InferType()(mod)
mod = preprocess.preprocess_ext_io()(mod)
return mod
def test_single_io():
"""
This test will test the pass wont touch external functions that
have a single input and a single output.
"""
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
mod = create_graph()
exp = create_graph()
mod = preprocess.preprocess_ext_io()(mod)
assert tvm.ir.structural_equal(mod, exp, map_free_vars=True)
def create_ethosu_partition(mod):
mod["main"] = EthosUAnnotator().visit(mod["main"])
mod = relay.transform.MergeCompilerRegions()(mod)
mod = relay.transform.InferType()(mod)
mod = relay.transform.PartitionGraph()(mod)
mod = relay.transform.InferType()(mod)
mod = preprocess.preprocess_ext_io()(mod)
return mod
def partition_ethosu_by_table(mod, pattern_table):
"""In case only the legalization part is supported for an operator, we don't
want to add the operator's pattern to the pattern table so that the compiler
wouldn't attempt to offload an operator without full stack support."""
mod = relay.transform.InferType()(mod)
mod = relay.transform.MergeComposite(pattern_table)(mod)
mod = relay.transform.AnnotateTarget("ethos-u")(mod)
mod = relay.transform.MergeCompilerRegions()(mod)
mod = relay.transform.InferType()(mod)
mod = relay.transform.PartitionGraph()(mod)
mod = relay.transform.InferType()(mod)
mod = preprocess.preprocess_ext_io()(mod)
return mod
def test_2ins_single_out():
"""
The test is check two inputs and a single output of external function
"""
def create_graph():
def create_external_func1(mod_, compiler_name, symbol_name):
x_int = relay.var("x_int", shape=(10, 10))
w0_int = relay.var("w0_int", shape=(10, 10))
z0 = relay.add(x_int, w0_int)
f1 = relay.Function([x_int, w0_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))
w0 = relay.var("w0", shape=(10, 10))
glb_symbol_f1, mod = create_external_func1(mod, "ethosu", "ethosu_0")
r = relay.Call(glb_symbol_f1, [x, w0])
main = relay.Function([x, w0], r)
mod["main"] = main
mod = relay.transform.InferType()(mod)
return mod
def expected():
def create_external_func1(mod_, compiler_name, symbol_name):
ifms_int = relay.var("ifms_int", shape=[200])
# splits
(x_int_flat, w0_int_flat) = relay.split(ifms_int, [100])
# reshapes
x_int = relay.reshape(x_int_flat, newshape=(10, 10))
w0_int = relay.reshape(w0_int_flat, newshape=(10, 10))
z0 = relay.add(x_int, w0_int)
f1 = relay.Function([ifms_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))
w0 = relay.var("w0", shape=(10, 10))
# reshapes
x_reshaped = relay.reshape(x, newshape=100)
w0_reshaped = relay.reshape(w0, newshape=100)
# concat
ifms = relay.concatenate((x_reshaped, w0_reshaped), 0)
glb_symbol_f1, mod = create_external_func1(mod, "ethosu", "ethosu_0")
r = relay.Call(glb_symbol_f1, [ifms])
main = relay.Function([x, w0], r)
mod["main"] = main
mod = relay.transform.InferType()(mod)
return mod
mod = create_graph()
exp = expected()
mod = preprocess.preprocess_ext_io()(mod)
assert tvm.ir.structural_equal(mod, exp, map_free_vars=True)
def test_4ins_2outs():
"""
The test is to check a 4 inputs and two outputs of external function.
This just stand as a general test for multiple ins/outs.
"""
def create_graph():
def create_external_func1(mod_, compiler_name, symbol_name):
x_int = relay.var("x_int", shape=(10, 10))
w0_int = relay.var("w0_int", shape=(10, 10))
w1_int = relay.var("w1_int", shape=(10, 10))
w2_int = relay.var("w2_int", shape=(10, 10))
z0 = relay.add(x_int, w0_int)
p0 = relay.subtract(z0, w1_int)
q0 = relay.multiply(z0, w2_int)
f1_o_tuple = relay.Tuple([p0, q0])
f1 = relay.Function([x_int, w0_int, w1_int, w2_int], f1_o_tuple)
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))
w0 = relay.var("w0", shape=(10, 10))
w1 = relay.var("w1", shape=(10, 10))
w2 = relay.var("w2", shape=(10, 10))
glb_symbol_f1, mod = create_external_func1(mod, "ethosu", "ethosu_0")
pq_tuple = relay.Call(glb_symbol_f1, [x, w0, w1, w2])
p0 = relay.TupleGetItem(pq_tuple, 0)
q0 = relay.TupleGetItem(pq_tuple, 1)
r = relay.concatenate((p0, q0), axis=0)
main = relay.Function([x, w0, w1, w2], r)
mod["main"] = main
mod = relay.transform.InferType()(mod)
return mod
def expected():
def create_external_func1(mod_, compiler_name, symbol_name):
ifms_int = relay.var("ifms_int", shape=[400])
# splits
(x_int_flat, w0_int_flat, w1_int_flat, w2_int_flat) = relay.split(
ifms_int, [100, 200, 300]
)
# reshapes
x_int = relay.reshape(x_int_flat, newshape=(10, 10))
w0_int = relay.reshape(w0_int_flat, newshape=(10, 10))
w1_int = relay.reshape(w1_int_flat, newshape=(10, 10))
w2_int = relay.reshape(w2_int_flat, newshape=(10, 10))
z0 = relay.add(x_int, w0_int)
p0 = relay.subtract(z0, w1_int)
q0 = relay.multiply(z0, w2_int)
# f1_o_tuple = relay.Tuple([p0, q0])
# 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([ifms_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))
w0 = relay.var("w0", shape=(10, 10))
w1 = relay.var("w1", shape=(10, 10))
w2 = relay.var("w2", shape=(10, 10))
# reshapes
x_reshaped = relay.reshape(x, newshape=100)
w0_reshaped = relay.reshape(w0, newshape=100)
w1_reshaped = relay.reshape(w1, newshape=100)
w2_reshaped = relay.reshape(w2, newshape=100)
# concat
ifms = relay.concatenate((x_reshaped, w0_reshaped, w1_reshaped, w2_reshaped), 0)
# call
glb_func, mod = create_external_func1(mod, "ethosu", "ethosu_0")
ofms = relay.Call(glb_func, [ifms])
# 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, w0, w1, w2], r)
mod["main"] = main
mod = relay.transform.InferType()(mod)
return mod
mod = create_graph()
exp = expected()
mod = preprocess.preprocess_ext_io()(mod)
assert tvm.ir.structural_equal(mod, exp, map_free_vars=True)
def test_single_in_2outs():
"""
The test is to check a single input and two outputs of external function
"""
def create_graph():
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)
f1_o_tuple = relay.Tuple([p0, q0])
f1 = relay.Function([x_int], f1_o_tuple)
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")
pq_tuple = relay.Call(glb_symbol_f1, [x])
p0 = relay.TupleGetItem(pq_tuple, 0)
q0 = relay.TupleGetItem(pq_tuple, 1)
r = relay.concatenate((p0, q0), axis=0)
main = relay.Function([x], r)
mod["main"] = main
mod = relay.transform.InferType()(mod)
return mod
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
mod = create_graph()
exp = expected()
mod = relay.transform.InferType()(mod)
mod = preprocess.preprocess_ext_io()(mod)
assert tvm.ir.structural_equal(mod, exp, map_free_vars=True)
