NicolaLancellotti commented on a change in pull request #9384:
URL: https://github.com/apache/tvm/pull/9384#discussion_r738596730
##########
File path: python/tvm/relay/backend/contrib/ethosu/legalize.py
##########
@@ -286,14 +288,125 @@ def callback(
@ir.transform.module_pass(opt_level=1)
-class LegalizeEthosUDepthwiseConv2D:
- """This is the pass that wraps the EthosUDepthwiseConv2DRewriter"""
+class LegalizeDepthwiseConv2D:
+ """This is the pass that wraps the DepthwiseConv2DRewriter"""
+
+ def transform_module(
+ self, mod: tvm.ir.IRModule, ctx: tvm.ir.transform.PassContext
+ ) -> tvm.ir.IRModule:
+ for global_var, func in mod.functions.items():
+ func = rewrite(DepthwiseConv2DRewriter(), func)
+ mod.update_func(global_var, func)
+ return mod
+
+ def __call__(self, *args, **kwargs):
+ pass
+
+
+class PoolingRewriter(DFPatternCallback):
+ """Convert ethosu.avgpool2d and ethosu.maxpool2d composite functions to
+ ethosu_pooling operators"""
+
+ def __init__(
+ self,
+ params_class: Type,
+ pattern: CallPattern,
+ ):
+ super().__init__(require_type=True)
+ self.params_class = params_class
+ self.pattern = pattern
+
+ def callback(
+ self, pre: tvm.relay.Expr, post: tvm.relay.Expr, node_map:
tvm.ir.container.Map
+ ) -> tvm.relay.Expr:
+ params = self.params_class(post.op.body)
+ params.ifm.tensor = post.args[0]
+ channels_map = {
+ "NHWC": 3,
+ }
+ if str(params.ofm.layout) not in channels_map.keys():
+ raise UnsupportedLayout(str(params.ofm.layout))
+
+ activation_map = {"clip": "CLIP"}
+ if params.activation:
+ activation = activation_map[params.activation.op.name]
+ clip_min = int(params.activation.attrs.a_min)
+ clip_max = int(params.activation.attrs.a_max)
+ else:
+ activation = "NONE"
+ clip_min = 0
+ clip_max = 0
+
+ # Activations requiring LUT is currently not supported, so setting it
to an empty list
+ lut = relay.const([], dtype="int8")
+
+ ethosu_pooling = ethosu_ops.ethosu_pooling(
Review comment:
No, I'll remove the variable.
--
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]
