tristan-arm commented on a change in pull request #8300:
URL: https://github.com/apache/tvm/pull/8300#discussion_r691187589
##########
File path: python/tvm/relay/op/contrib/ethosn.py
##########
@@ -61,6 +61,53 @@ def partition_for_ethosn(mod, params=None, **opts):
-------
ret : annotated and partitioned module.
"""
+ if opts:
+ tops = opts.get("tops", None)
+ ple_ratio = opts.get("ple_ratio", None)
+ sram_size = opts.get("sram_size", None)
+ if tops or ple_ratio or sram_size:
+ raise ValueError(
+ "Setting tops, ple_ratio or sram_size has no effect when
targeting Ethos-N77"
+ )
+
+ if params:
+ mod["main"] = bind_params_by_name(mod["main"], params)
+
+ seq = tvm.transform.Sequential(
+ [
+ transform.InferType(),
+ transform.MergeComposite(pattern_table()),
+ transform.AnnotateTarget("ethos-n"),
+ transform.MergeCompilerRegions(),
+ transform.PartitionGraph(),
+ ]
+ )
+
+ return seq(mod), None
+
+
+def partition_for_ethosn78(mod, params=None, **opts):
+ """Partition the graph greedily offloading supported
+ operators to Arm Ethos-N NPU.
+
+ Parameters
+ ----------
+ mod : Module
+ The module to run passes on.
+ params : Optional[Dict[str, NDArray]]
+ Constant input parameters.
+
+ Returns
+ -------
+ ret : annotated and partitioned module.
+ """
+ config = opts
+ if not config:
+ config = {}
+
+ if config.get("variant", "").lower() != "ethos-n78":
Review comment:
I've used config.get("variant", "") here to handle the variant undefined
scenario.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
