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 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)