ekalda commented on code in PR #14765:
URL: https://github.com/apache/tvm/pull/14765#discussion_r1195229954
##########
python/tvm/relay/backend/contrib/ethosu/legalize.py:
##########
@@ -1447,6 +1447,84 @@ def callback(
)
+class ChannelPadRewriter(DFPatternCallback):
+ """Convert ethos-u.pad2d composite function to the Relay concatenate
operation"""
Review Comment:
```suggestion
"""Convert ethos-u.channel-pad composite function to the Relay
concatenate operation"""
```
##########
python/tvm/relay/backend/contrib/ethosu/legalize.py:
##########
@@ -1447,6 +1447,84 @@ def callback(
)
+class ChannelPadRewriter(DFPatternCallback):
+ """Convert ethos-u.pad2d composite function to the Relay concatenate
operation"""
+
+ def __init__(self):
+ super().__init__(require_type=True)
+ self.pattern = (
+ wildcard().has_attr({"Composite":
ethosu_patterns.ChannelPadParams.composite_name})
+ )(wildcard())
+
+ def callback(
+ self, pre: tvm.relay.Expr, post: tvm.relay.Expr, node_map:
tvm.ir.container.Map
+ ) -> tvm.relay.Expr:
+ params = ethosu_patterns.ChannelPadParams(post.op.body)
+ params.ifm.tensor = post.args[0]
+
+ concat_args = list()
+ # Activations requiring LUT is currently not supported, so setting it
to an empty list
Review Comment:
I know every operator here has this copy pasted legacy comment, but let's
remove it... Firstly LUT based activations are supported and secondly it could
leave an impression that implementing something like fused pad + sigmoid is a
TODO.
##########
tests/python/contrib/test_ethosu/test_legalize.py:
##########
@@ -462,6 +463,118 @@ def verify(ext_func):
verify(mod["tvmgen_default_ethos_u_main_0"])
+def test_tflite_conv2d_with_separate_channel_padding_legalize():
+ dtype = "int8"
+ ifm_shape = (1, 55, 34, 3)
+ kernel_shape = (3, 2)
+ strides = (1, 1)
+ dilation = (2, 1)
+ padding_ch = (1, 1)
+
+ class ArePadOnGraph(ExprVisitor):
+ """
+ Visits the Graph recursively and checks if it contains 'nn.pad' op
+ """
+
+ def __init__(self):
+ ExprVisitor.__init__(self)
+ self.on_graph = False
+
+ def visit_call(self, call):
+ if isinstance(call.op, tvm.ir.Op):
+ if str(call.op.name) == "nn.pad":
+ self.on_graph = True
+
+ return super().visit_call(call)
+
+ def are_pad_on_graph(self, subgraph) -> bool:
+ """
+ This function recursively visits the graph and checks if 'nn.pad'
op is ongraph
Review Comment:
Nit:
```suggestion
This function recursively visits the graph and checks if
'nn.pad' op is on graph
```
##########
python/tvm/relay/backend/contrib/ethosu/legalize.py:
##########
@@ -1447,6 +1447,84 @@ def callback(
)
+class ChannelPadRewriter(DFPatternCallback):
+ """Convert ethos-u.pad2d composite function to the Relay concatenate
operation"""
+
+ def __init__(self):
+ super().__init__(require_type=True)
+ self.pattern = (
+ wildcard().has_attr({"Composite":
ethosu_patterns.ChannelPadParams.composite_name})
+ )(wildcard())
+
+ def callback(
+ self, pre: tvm.relay.Expr, post: tvm.relay.Expr, node_map:
tvm.ir.container.Map
+ ) -> tvm.relay.Expr:
+ params = ethosu_patterns.ChannelPadParams(post.op.body)
+ params.ifm.tensor = post.args[0]
+
+ concat_args = list()
+ # Activations requiring LUT is currently not supported, so setting it
to an empty list
+ lut = relay.const([], dtype="int8")
+ # pad channels before
+ if params.ch_padding[0] > 0:
+ shape1 = list(params.ifm.shape)
+ shape1[3] = params.ch_padding[0].value
+ pad_channels = relay.Constant(
+ tvm.nd.array(
+ np.full(
+ shape=shape1,
+ fill_value=int(params.ifm.q_params.zero_point),
+ dtype=params.ifm.dtype,
+ )
+ )
+ )
+ identity1 = ethosu_ops.ethosu_identity(
+ ifm=pad_channels,
+ lut=lut,
+ ifm_scale=float(params.ifm.q_params.scale_f32),
+ ifm_zero_point=int(params.ifm.q_params.zero_point),
+ ofm_scale=float(params.ofm.q_params.scale_f32),
+ ofm_zero_point=int(params.ofm.q_params.zero_point),
+ )
+ concat_args.append(identity1)
+
+ identity2 = ethosu_ops.ethosu_identity(
+ ifm=params.ifm.tensor,
+ lut=lut,
+ ifm_scale=float(params.ifm.q_params.scale_f32),
+ ifm_zero_point=int(params.ifm.q_params.zero_point),
+ ofm_scale=float(params.ofm.q_params.scale_f32),
+ ofm_zero_point=int(params.ofm.q_params.zero_point),
+ )
+ concat_args.append(identity2)
+
+ # pad channels after
+ if params.ch_padding[1] > 0:
+ shape3 = list(params.ifm.shape)
+ shape3[3] = params.ch_padding[1].value
+ pad_channels3 = relay.Constant(
+ tvm.nd.array(
+ np.full(
+ shape=shape3,
+ fill_value=int(params.ifm.q_params.zero_point),
+ dtype=params.ifm.dtype,
+ )
+ )
+ )
+ identity3 = ethosu_ops.ethosu_identity(
+ ifm=pad_channels3,
+ lut=lut,
+ ifm_scale=float(params.ifm.q_params.scale_f32),
+ ifm_zero_point=int(params.ifm.q_params.zero_point),
+ ofm_scale=float(params.ofm.q_params.scale_f32),
+ ofm_zero_point=int(params.ofm.q_params.zero_point),
+ )
+ concat_args.append(identity3)
+
+ axis = 3
+ return relay.op.concatenate(relay.Tuple(concat_args), axis=axis)
Review Comment:
Since it is not used elsewhere, maybe just
```suggestion
return relay.op.concatenate(relay.Tuple(concat_args), axis=3)
```
##########
tests/python/contrib/test_ethosu/test_legalize.py:
##########
@@ -760,7 +873,98 @@ def verify(ext_func):
ethosu.PadParams.composite_name,
ethosu.pad_pattern(),
lambda pat: ethosu.PadParams(pat).is_valid(),
+ ),
+ ]
+
+ tflite_graph = create_tflite_graph()
+ tflite_model = tflite.Model.Model.GetRootAsModel(tflite_graph, 0)
+
+ mod, params = relay.frontend.from_tflite(
+ tflite_model,
+ shape_dict={"input": ifm_shape},
+ dtype_dict={"input": dtype},
+ )
+
+ mod["main"] = bind_params_by_name(mod["main"], params)
+ mod = partition_ethosu_by_table(mod, pad_pattern_table)
+
+ mod["tvmgen_default_ethos_u_main_0"] = dataflow_pattern.rewrite(
+ legalize.PadRewriter(), mod["tvmgen_default_ethos_u_main_0"]
+ )
+ verify(mod["tvmgen_default_ethos_u_main_0"])
+
+
[email protected]("ifm_shape", [(1, 55, 55, 3), (1, 23, 32, 7)])
[email protected]("channel_padding", [(0, 1), (1, 1), (5, 2)])
[email protected]("const_value", [0, 5, 125, -5])
+def test_tflite_separate_channel_padding_legalize(ifm_shape, channel_padding,
const_value):
Review Comment:
Shouldn't this test be using `ChannelPadRewriter` and then check in `verify`
that the concatenates got created?
##########
tests/python/contrib/test_ethosu/test_legalize.py:
##########
@@ -462,6 +463,118 @@ def verify(ext_func):
verify(mod["tvmgen_default_ethos_u_main_0"])
+def test_tflite_conv2d_with_separate_channel_padding_legalize():
+ dtype = "int8"
+ ifm_shape = (1, 55, 34, 3)
+ kernel_shape = (3, 2)
+ strides = (1, 1)
+ dilation = (2, 1)
+ padding_ch = (1, 1)
+
+ class ArePadOnGraph(ExprVisitor):
+ """
+ Visits the Graph recursively and checks if it contains 'nn.pad' op
+ """
+
+ def __init__(self):
+ ExprVisitor.__init__(self)
+ self.on_graph = False
+
+ def visit_call(self, call):
+ if isinstance(call.op, tvm.ir.Op):
+ if str(call.op.name) == "nn.pad":
+ self.on_graph = True
+
+ return super().visit_call(call)
+
+ def are_pad_on_graph(self, subgraph) -> bool:
+ """
+ This function recursively visits the graph and checks if 'nn.pad'
op is ongraph
+ """
+ self.visit(subgraph)
+ return self.on_graph
+
+ def create_tflite_graph_single():
+ class Model(tf.Module):
+ @tf.function
+ def tf_function(self, x):
+ tf_strides = [1, strides[0], strides[1], 1]
+ op = tf.pad(
+ x,
+ [[0, 0], [0, 0], [0, 0], [padding_ch[0], padding_ch[1]]],
+ "CONSTANT",
+ )
+ # HWIO
+ weight_shape = [
+ kernel_shape[0],
+ kernel_shape[1],
+ ifm_shape[3] + padding_ch[0] + padding_ch[1],
+ 3,
+ ]
+ weight = tf.constant(np.random.uniform(size=weight_shape),
dtype=tf.float32)
+ return tf.nn.conv2d(
+ op,
+ weight,
+ strides=tf_strides,
+ padding="VALID",
+ dilations=dilation,
+ )
+
+ model = Model()
+ concrete_func = model.tf_function.get_concrete_function(
+ tf.TensorSpec(ifm_shape, dtype=tf.float32)
+ )
+ # Convert the model
+ def representative_dataset():
+ for _ in range(100):
+ data = np.random.rand(*tuple(ifm_shape))
+ yield [data.astype(np.float32)]
+
+ converter =
tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
+ converter.optimizations = [tf.lite.Optimize.DEFAULT]
+ converter.representative_dataset = representative_dataset
+ converter.target_spec.supported_ops =
[tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
+ converter.inference_input_type = tf.int8
+ converter.inference_output_type = tf.int8
+ tflite_model = converter.convert()
+ return tflite_model
+
+ def verify(ext_func):
+
+ assert ArePadOnGraph().are_pad_on_graph(ext_func.body) == True
+
+ conv2d_pattern_table = [
+ (
+ ethosu.ChannelPadParams.composite_name,
+ ethosu.pad_pattern(),
+ lambda pat: ethosu.ChannelPadParams(pat).is_valid(),
+ ),
+ (
+ ethosu.QnnConv2DParams.composite_name,
+ ethosu.qnn_conv2d_pattern(),
+ lambda pat: ethosu.QnnConv2DParams(pat).is_valid(),
+ ),
+ ]
+
+ tflite_graph = create_tflite_graph_single()
+ tflite_model = tflite.Model.Model.GetRootAsModel(tflite_graph, 0)
+
+ mod, conv_params = relay.frontend.from_tflite(
+ tflite_model,
+ shape_dict={"input": ifm_shape},
+ dtype_dict={"input": dtype},
+ )
+
+ mod["main"] = bind_params_by_name(mod["main"], conv_params)
+ mod = partition_ethosu_by_table(mod, conv2d_pattern_table)
+
+ mod["tvmgen_default_ethos_u_main_0"] = dataflow_pattern.rewrite(
+ legalize.Conv2DRewriter(), mod["tvmgen_default_ethos_u_main_0"]
+ )
+
+ verify(mod["tvmgen_default_ethos_u_main_0"])
Review Comment:
I'm a bit confused about what that test does... It creates a TFLite graph
with channel pad and conv2d, then partitions them for microNPU, then legalizes
the conv2d into `ethosu.conv2d` and then checks that the Relay pad is still
there?
--
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]
