def primfunc_to_artifact(
primfunc: tvm.tir.PrimFunc) -> util.CompilationArtifact:
"""
This is the hook for python-based lowering of TIR PrimFunc
that has undergone unified optimization to compilation
artifact destined for the microNPU.
Parameters
----------
primfunc : tir.PrimFunc
TIR PrimFunc that has undergone unified optimizations
Returns
-------
CompilationArtifact
This is a structure that holds the binary artifacts
for the microNPU
"""
symbol = str(primfunc.attrs["global_symbol"])
const_dict = primfunc.attrs["ethos-u.constants"]
tir_mod = tvm.IRModule()
tir_mod[symbol] = primfunc
const_dict_np = dict()
for buffer_var in const_dict.keys():
const_dict_np[buffer_var] = const_dict[buffer_var].numpy()
cmms, encoded_constants, base_addresses = tir_to_cs_translator.translate(
tir_mod, const_dict_np)
return util.CompilationArtifact(symbol, cmms, encoded_constants,
base_addresses)
def _compile(ext_func):
"""
This is the main wrapper that accepts an external
relay function and runs all the passes to lower it down
to command stream
Parameters
----------
ext_func : tvm.relay.function.Function
The partitioned relay function
Returns
-------
cs : str
An hex string of the bytes of command stream
encoded_constants : str
An hex string of the bytes that includes concat'd
encoded weights, encoded biases and scales.
scratch_size : int
The size of the scratch buffer needed.
"""
mod = tvm.IRModule()
mod["main"] = ext_func
mod = LegalizeEthosU()(mod)
mod = relay.transform.InferType()(mod)
# We are currently using copy_constants scheduler In the long run,
# this should be a single intelligent and a composite scheduler
# that can perform scheduling based on user inputs such as
# scratch memory size.
tir_mod, params = lower_to_tir(mod["main"], copy_constants())
cmms, encoded_constants, scratch_size = tir_to_cs_translator.translate(
tir_mod, params)
return cmms, encoded_constants, scratch_size
def primfunc_to_artifact(
primfunc: tvm.tir.PrimFunc) -> util.CompilationArtifact:
"""
This is the hook for python-based lowering of TIR PrimFunc
that has undergone unified optimization to compilation
artifact destined for the microNPU.
Parameters
----------
primfunc : tir.PrimFunc
TIR PrimFunc that has undergone unified optimizations
Returns
-------
CompilationArtifact
This is a structure that holds the binary artifacts
for the microNPU
"""
symbol = str(primfunc.attrs["global_symbol"])
const_dict = primfunc.attrs["ethos-u.constants"]
input_size = primfunc.attrs["ethos-u.input_size"]
output_size = primfunc.attrs["ethos-u.output_size"]
tir_mod = tvm.IRModule()
tir_mod[symbol] = primfunc
const_dict_with_int_keys = dict()
for idx in const_dict.keys():
const_dict_with_int_keys[int(idx)] = const_dict[idx].numpy()
cmms, encoded_constants, scratch_size = tir_to_cs_translator.translate(
tir_mod, const_dict_with_int_keys)
return util.CompilationArtifact(cmms, encoded_constants, scratch_size,
input_size, output_size, symbol)
def test_constant_as_input():
"""Test to check that constants specified as inputs aren't
interpreted as an encoded constant."""
def get_graph():
dtype = "uint8"
ifm = relay.var("ifm", shape=(1, 16, 16, 32), dtype=dtype)
conv1 = make_ethosu_conv2d(
ifm,
32,
16,
(1, 1),
(0, 0),
(1, 1),
(1, 1),
)
scalar = relay.const(np.ones((1, 1, 1, 1), dtype=dtype), dtype=dtype)
add1 = make_ethosu_binary_elementwise(conv1,
scalar,
ifm_channels=32,
ifm2_channels=1,
operator_type="ADD",
ofm_dtype=dtype)
func = relay.Function(relay.analysis.free_vars(add1), add1)
func = run_opt_pass(func, relay.transform.InferType())
return func
tir_mod, params = _lower_to_tir(get_graph(), copy_constants())
# Check tile address for the scalar constant input hasn't been
# overwritten.
extern_calls = tir_mod["main"].body.body.body.body.body
binary_elementwise = extern_calls[-1].value
args = binary_elementwise.args
reason = "Tile address overwritten"
assert args[26] == 0, reason
assert args[27] == 0, reason
assert args[28] == 0, reason
# More generally, check compiles successfully to make sure
# nothing else was overrwritten.
# With Target Hooks the TIR module needs a target attached
# and lowered via make unpacked API.
tir_mod["main"] = tir_mod["main"].with_attr("target",
tvm.target.Target("ethos-u"))
tir_mod = tvm.tir.transform.MakeUnpackedAPI()(tir_mod)
tir_to_cs_translator.translate(tir_mod, params)