def build(mod, params, npu=True, expected_host_ops=0, npu_partitions=1):
relay.backend.compile_engine.get().clear()
with tvm.transform.PassContext(
opt_level=3, config={"relay.ext.ethos-n.options": {
"variant": 0
}}):
with tvm.target.Target("llvm"):
if npu:
f = relay.build_module.bind_params_by_name(mod["main"], params)
mod = tvm.IRModule()
mod["main"] = f
pattern = get_pattern_table("ethos-n")
mod = relay.transform.MergeComposite(pattern)(mod)
mod = relay.transform.AnnotateTarget("ethos-n")(mod)
mod = relay.transform.MergeCompilerRegions()(mod)
mod = relay.transform.PartitionGraph()(mod)
host_op_count = get_host_op_count(mod)
assert (host_op_count == expected_host_ops
), "Got {} host operators, expected {}".format(
host_op_count, expected_host_ops)
partition_count = 0
for global_var in mod.get_global_vars():
if "ethos-n" in global_var.name_hint:
partition_count += 1
assert (npu_partitions == partition_count
), "Got {} ethos-n partitions, expected {}".format(
partition_count, npu_partitions)
return relay.build(mod, params=params)
def test_reshape_failure(input_shape, output_shape, dtype, err_msg):
np.random.seed(0)
model, params = _get_model(input_shape, output_shape, dtype)
mod = tei.make_module(model, params)
pattern = get_pattern_table("ethos-n")
mod = tei.make_module(model, params)
mod = relay.transform.MergeComposite(pattern)(mod)
mod = tei.make_ethosn_partition(mod["main"].body)
tei.test_error(mod, {}, err_msg)
def build(mod, params, npu=True, expected_host_ops=0, npu_partitions=1):
"""Build a network with or without Ethos-N offloading.
Parameters
----------
mod : IRModule
The Relay module to build.
params : dict of str to NDArray
The weights to build with.
npu : bool, optional
Whether to build with Ethos-N offloading.
expected_host_ops : int, optional
The number of ops expected to remain on the host.
npu_partitions : int, optional
The number of Ethos-N partitions expected.
"""
relay.backend.te_compiler.get().clear()
with tvm.transform.PassContext(
opt_level=3, config={"relay.ext.ethos-n.options": {"variant": get_ethosn_variant()}}
):
with tvm.target.Target("llvm"):
if npu:
f = relay.build_module.bind_params_by_name(mod["main"], params)
mod = tvm.IRModule()
mod["main"] = f
pattern = get_pattern_table("ethos-n")
mod = relay.transform.InferType()(mod)
mod = relay.transform.MergeComposite(pattern)(mod)
mod = relay.transform.AnnotateTarget("ethos-n")(mod)
mod = relay.transform.InferType()(mod)
mod = relay.transform.MergeCompilerRegions()(mod)
mod = relay.transform.InferType()(mod)
mod = relay.transform.PartitionGraph()(mod)
host_op_count = get_host_op_count(mod)
assert (
host_op_count == expected_host_ops
), "Got {} host operators, expected {}".format(host_op_count, expected_host_ops)
attrs = [
mod[var.name_hint].attrs
for var in mod.get_global_vars()
if mod[var.name_hint].attrs
]
partition_count = sum(
[
key == "Compiler" and value == "ethos-n"
for attr in attrs
for key, value in attr.items()
]
)
assert (
npu_partitions == partition_count
), "Got {} ethos-n partitions, expected {}".format(partition_count, npu_partitions)
return relay.build(mod, params=params)
def test_reshape_failure():
trials = [
(
(1, 15, 4, 1),
(1, 15, -2),
"uint8",
"reshape dimension=-2, reshape dimension must be >= -1",
),
]
np.random.seed(0)
for input_shape, output_shape, dtype, err_msg in trials:
model, params = _get_model(input_shape, output_shape, dtype)
mod = tei.make_module(model, params)
pattern = get_pattern_table("ethos-n")
mod = tei.make_module(model, params)
mod = relay.transform.MergeComposite(pattern)(mod)
mod = tei.make_ethosn_partition(mod["main"].body)
tei.test_error(mod, {}, err_msg)
