This is an automated email from the ASF dual-hosted git repository.
ekalda pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new d9c1ba60f4 [microNPU][ETHOSU] Add offloading to the NPU the
nn.avg_pool2d operator with a stride > 3 (#14861)
d9c1ba60f4 is described below
commit d9c1ba60f446b0363a2ff3b65b63481898ec2221
Author: sergio-grovety <[email protected]>
AuthorDate: Wed May 24 16:54:58 2023 +0300
[microNPU][ETHOSU] Add offloading to the NPU the nn.avg_pool2d operator
with a stride > 3 (#14861)
The nn.avg_pool2d operator with a stride size greater than 3 in any of the
spatial dimensions is rewritten as ethosu avg_pool with strides=[1,1] if the
case satisfies the additional conditions (no AvgPool2D padding, spatial
dimensions of ifm and shape of pooling are equal).
---------
Co-authored-by: Sergey Smirnov
<[email protected]>
Co-authored-by: arina-grovety <>
Co-authored-by: Arina Naumova (grovety.com) <[email protected]>
Co-authored-by: Arina <[email protected]>
---
.../tvm/relay/backend/contrib/ethosu/legalize.py | 11 +-
python/tvm/relay/op/contrib/ethosu.py | 27 ++-
tests/python/contrib/test_ethosu/test_codegen.py | 32 ++++
tests/python/contrib/test_ethosu/test_legalize.py | 211 +++++++++++++++++++++
4 files changed, 277 insertions(+), 4 deletions(-)
diff --git a/python/tvm/relay/backend/contrib/ethosu/legalize.py
b/python/tvm/relay/backend/contrib/ethosu/legalize.py
index b4e8124d14..175f97c87e 100644
--- a/python/tvm/relay/backend/contrib/ethosu/legalize.py
+++ b/python/tvm/relay/backend/contrib/ethosu/legalize.py
@@ -618,6 +618,15 @@ class PoolingRewriter(DFPatternCallback):
# Activations requiring LUT is currently not supported, so setting it
to an empty list
lut = relay.const([], dtype="int8")
+ # If ethosu.avgpool2d has strides which are not supported by the NPU,
convert
+ # ethosu.avgpool2d composite functions to ethosu_pooling operator with
stride=[1, 1].
+ # Since the spatial dimensions of ifm and the pooling kernel coincide
and the padding
+ # is [0, 0, 0, 0], the application of the pooling kernel will be done
only once,
+ # which will give us the desired output
+ strides = params.strides
+ if params.strides[0] > 3 or params.strides[1] > 3:
+ strides = [1, 1]
+
return ethosu_ops.ethosu_pooling(
ifm=post.args[0],
lut=lut,
@@ -629,7 +638,7 @@ class PoolingRewriter(DFPatternCallback):
pool_shape=params.pool_shape,
ofm_channels=params.ofm.shape[channels_map[str(params.ofm.layout)]],
ofm_dtype=params.ofm.dtype,
- strides=params.strides,
+ strides=strides,
padding=params.padding,
activation=activation,
clip_min=clip_min,
diff --git a/python/tvm/relay/op/contrib/ethosu.py
b/python/tvm/relay/op/contrib/ethosu.py
index 71b419507b..0796ccf62a 100644
--- a/python/tvm/relay/op/contrib/ethosu.py
+++ b/python/tvm/relay/op/contrib/ethosu.py
@@ -101,6 +101,23 @@ def check_strides(strides: List[int], stride_range=None)
-> bool:
return True
+def check_same_ifm_and_kernel_shape(padding, ifm_shape, pool_shape):
+ """
+ This function checks whether AvgPool2D or MaxPool2D could be legalized as
ethosu_pooling
+ supported by the NPU.
+ We consider only specific case: when there is no AvgPool2D padding, the
spatial
+ dimensions of ifm and the shape of pooling are equal, but stride size
exceed 3
+ by any of dimensions, e.g:
+ ifm: (1, 8, 8, _), strides: (8, 8), pool_shape: (8, 8)
+ ifm: (1, 25, 5, _), strides: (25, 5), pool_shape: (25, 5)
+ """
+ if list(padding) != [0, 0, 0, 0]:
+ return False
+ if [ifm_shape[1], ifm_shape[2]] != list(pool_shape):
+ return False
+ return True
+
+
def check_valid_dtypes(tensor_params: List[TensorParams], supported_dtypes:
List[type]) -> bool:
"""This function checks whether dtypes are supported by the NPU"""
for tep in tensor_params:
@@ -595,7 +612,9 @@ class MaxPool2DParams:
return False
if self.ifm.dtype != self.ofm.dtype:
return False
- if not check_strides(self.strides):
+ if not check_strides(self.strides) and not
check_same_ifm_and_kernel_shape(
+ self.padding, self.ifm.shape, self.pool_shape
+ ):
return False
if not check_batch_size(self.ifm):
return False
@@ -655,7 +674,9 @@ class AvgPool2DParams:
return False
if self.ifm.dtype != self.ofm.dtype:
return False
- if not check_strides(self.strides):
+ if not check_strides(self.strides) and not
check_same_ifm_and_kernel_shape(
+ self.padding, self.ifm.shape, self.pool_shape
+ ):
return False
if not check_batch_size(self.ifm):
return False
@@ -665,7 +686,7 @@ class AvgPool2DParams:
return False
if not check_pool_shape(self.pool_shape):
return False
- # Averge pool with padding only supports 1 <= pool_shape <= 8
+ # Average pool with padding only supports 1 <= pool_shape <= 8
if list(self.padding) != [0, 0, 0, 0] and (
self.pool_shape[0] > 8 or self.pool_shape[1] > 8
):
diff --git a/tests/python/contrib/test_ethosu/test_codegen.py
b/tests/python/contrib/test_ethosu/test_codegen.py
index ef91b75efa..cb1592c041 100644
--- a/tests/python/contrib/test_ethosu/test_codegen.py
+++ b/tests/python/contrib/test_ethosu/test_codegen.py
@@ -338,6 +338,38 @@ def test_ethosu_pooling(
infra.compare_tvm_with_tflite(pooling, [ifm_shape], accel_type)
[email protected](
+ "accel_type",
+ ACCEL_TYPES,
+)
[email protected]("pooling_type", ["MAX", "AVG"])
[email protected](
+ "ifm_shape, pool_shape, strides, activation_function, padding",
+ [
+ ([1, 4, 4, 3], [4, 4], [4, 4], "NONE", "SAME"),
+ ([1, 4, 4, 3], [4, 4], [4, 4], "RELU", "VALID"),
+ ([1, 25, 5, 64], [25, 5], [25, 5], "NONE", "VALID"),
+ ([1, 25, 5, 64], [25, 5], [25, 5], "RELU", "SAME"),
+ ],
+)
+def test_ethosu_pooling_same_ifm_and_kernel_shape(
+ accel_type, pooling_type, ifm_shape, pool_shape, strides,
activation_function, padding
+):
+ np.random.seed(0)
+
+ @tf.function
+ def pooling(x):
+ if pooling_type == "MAX":
+ op = tf.nn.max_pool(x, pool_shape, strides, padding)
+ elif pooling_type == "AVG":
+ op = tf.nn.avg_pool(x, pool_shape, strides, padding)
+ if activation_function == "RELU":
+ op = tf.nn.relu(op)
+ return op
+
+ infra.compare_tvm_with_tflite(pooling, [ifm_shape], accel_type)
+
+
@pytest.mark.parametrize(
"accel_type",
["ethos-u55-256", "ethos-u65-256"],
diff --git a/tests/python/contrib/test_ethosu/test_legalize.py
b/tests/python/contrib/test_ethosu/test_legalize.py
index f87b2da983..1b643f8157 100644
--- a/tests/python/contrib/test_ethosu/test_legalize.py
+++ b/tests/python/contrib/test_ethosu/test_legalize.py
@@ -1114,6 +1114,217 @@ def test_tflite_pool2d_legalize(
verify(mod["tvmgen_default_ethos_u_main_0"])
[email protected]("pooling_type", ["MAX", "AVG"])
[email protected](
+ "ifm_shape, pool_shape, strides, activation_function, padding",
+ [
+ ([1, 4, 4, 3], [4, 4], [4, 4], "NONE", "SAME"),
+ ([1, 4, 4, 3], [4, 4], [4, 4], "RELU", "VALID"),
+ ([1, 25, 5, 64], [25, 5], [25, 5], "NONE", "VALID"),
+ ([1, 25, 5, 64], [25, 5], [25, 5], "RELU", "SAME"),
+ ],
+)
+def test_tflite_pool2d_same_ifm_and_kernel_shape_legalize(
+ pooling_type, ifm_shape, pool_shape, strides, activation_function, padding
+):
+ dtype = "int8"
+ strides_legalized = [1, 1]
+
+ def create_tflite_graph():
+ class Model(tf.Module):
+ @tf.function
+ def tf_function(self, x):
+ if pooling_type == "MAX":
+ op = tf.nn.max_pool(x, pool_shape, strides, padding)
+ elif pooling_type == "AVG":
+ op = tf.nn.avg_pool(x, pool_shape, strides, padding)
+ if activation_function == "RELU":
+ op = tf.nn.relu(op)
+ return op
+
+ 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 expected_mod():
+
+ expected_ir_string = ""
+
+ if activation_function == "NONE" and pooling_type == "AVG":
+ expected_ir_string = f"""
+ #[version = "0.0.5"]
+ def @main(%x: Tensor[{str(tuple(ifm_shape))}, {dtype}],
output_tensor_names=\
+ ["Identity"]) -> Tensor[(1, 1, 1, {str(ifm_shape[3])}),
{dtype}] {{
+ @tvmgen_default_ethos_u_main_0(%x)
+ }}
+
+ def @tvmgen_default_ethos_u_main_0(%y:
Tensor[{str(tuple(ifm_shape))}, {dtype}], \
+ Compiler="ethos-u", Primitive=1, Inline=1, \
+ global_symbol="tvmgen_default_ethos_u_main_0") ->
Tensor[(1, 1, 1, \
+ {str(ifm_shape[3])}), {dtype}] {{
+ %2 = fn (%z: Tensor[{str(tuple(ifm_shape))}, {dtype}], \
+ PartitionedFromPattern="cast_nn.avg_pool2d_cast_", \
+ Composite="ethos-u.avgpool2d") -> Tensor[(1, 1, 1,
{str(ifm_shape[3])}), \
+ {dtype}] {{
+ %0 = cast(%z, dtype="int32") ;
+ %1 = nn.avg_pool2d(%0, pool_size={str(pool_shape)},
strides={str(strides)}, \
+ padding=[0, 0, 0, 0], layout="NHWC") ;
+ cast(%1, dtype="{dtype}")
+ }} ;
+ %2(%y)
+ }}
+ """
+
+ if activation_function == "RELU" and pooling_type == "AVG":
+ expected_ir_string = f"""
+ #[version = "0.0.5"]
+ def @main(%x: Tensor[{str(tuple(ifm_shape))}, {dtype}],
output_tensor_names=\
+ ["Identity"]) -> Tensor[(1, 1, 1, {str(ifm_shape[3])}),
{dtype}] {{
+ @tvmgen_default_ethos_u_main_0(%x)
+ }}
+
+ def @tvmgen_default_ethos_u_main_0(%y:
Tensor[{str(tuple(ifm_shape))}, {dtype}], \
+ Compiler="ethos-u", Primitive=1, Inline=1, \
+ global_symbol="tvmgen_default_ethos_u_main_0") ->
Tensor[(1, 1, 1, \
+ {str(ifm_shape[3])}), {dtype}] {{
+ %3 = fn (%z: Tensor[{str(tuple(ifm_shape))}, {dtype}], \
+ PartitionedFromPattern="cast_nn.avg_pool2d_cast_clip_", \
+ Composite="ethos-u.avgpool2d") -> Tensor[(1, 1, 1,
{str(ifm_shape[3])}), \
+ {dtype}] {{
+ %0 = cast(%z, dtype="int32") ;
+ %1 = nn.avg_pool2d(%0, pool_size={str(pool_shape)},
strides={str(strides)}, \
+ padding=[0, 0, 0, 0], layout="NHWC") ;
+ %2 = cast(%1, dtype="{dtype}") ;
+ clip(%2, a_min=-128f, a_max=127f)
+ }} ;
+ %3(%y)
+ }}
+ """
+
+ if activation_function == "NONE" and pooling_type == "MAX":
+ expected_ir_string = f"""
+ #[version = "0.0.5"]
+ def @main(%x: Tensor[{str(tuple(ifm_shape))}, {dtype}],
output_tensor_names=\
+ ["Identity"]) -> Tensor[(1, 1, 1, {str(ifm_shape[3])}),
{dtype}] {{
+ @tvmgen_default_ethos_u_main_0(%x)
+ }}
+
+ def @tvmgen_default_ethos_u_main_0(%y:
Tensor[{str(tuple(ifm_shape))}, {dtype}], \
+ Compiler="ethos-u", Primitive=1, Inline=1, \
+ global_symbol="tvmgen_default_ethos_u_main_0") ->
Tensor[(1, 1, 1, \
+ {str(ifm_shape[3])}), {dtype}] {{
+ %0 = fn (%z: Tensor[{str(tuple(ifm_shape))}, {dtype}], \
+ PartitionedFromPattern="nn.max_pool2d_", \
+ Composite="ethos-u.maxpool2d") -> Tensor[(1, 1, 1,
{str(ifm_shape[3])}), \
+ {dtype}] {{
+ nn.max_pool2d(%z, pool_size={str(pool_shape)},
strides={str(strides)}, \
+ padding=[0, 0, 0, 0], layout="NHWC")
+ }} ;
+ %0(%y)
+ }}
+ """
+
+ if activation_function == "RELU" and pooling_type == "MAX":
+ expected_ir_string = f"""
+ #[version = "0.0.5"]
+ def @main(%x: Tensor[{str(tuple(ifm_shape))}, {dtype}] ,
output_tensor_names=\
+ ["Identity"]) -> Tensor[(1, 1, 1, {str(ifm_shape[3])}),
{dtype}] {{
+ @tvmgen_default_ethos_u_main_0(%x)
+ }}
+
+ def @tvmgen_default_ethos_u_main_0(%y:
Tensor[{str(tuple(ifm_shape))}, {dtype}] , \
+ Compiler="ethos-u", Primitive=1, Inline=1, \
+ global_symbol="tvmgen_default_ethos_u_main_0") ->
Tensor[(1, 1, 1, \
+ {str(ifm_shape[3])}), {dtype}] {{
+ %1 = fn (%z: Tensor[{str(tuple(ifm_shape))}, {dtype}] , \
+ PartitionedFromPattern="nn.max_pool2d_clip_", \
+ Composite="ethos-u.maxpool2d") -> Tensor[(1, 1, 1,
{str(ifm_shape[3])}), \
+ {dtype}] {{
+ %0 = nn.max_pool2d(%z, pool_size={str(pool_shape)},
strides={str(strides)}, \
+ padding=[0, 0, 0, 0], layout="NHWC");
+ clip(%0, a_min=-128f, a_max=127f)
+ }};
+ %1(%y)
+ }}
+ """
+
+ return tvm.relay.fromtext(expected_ir_string)
+
+ def verify(ext_func):
+ ofm_shape = infra.compute_ofm_shape(ifm_shape, padding, pool_shape,
strides)
+ op = ext_func.body
+ assert list(op.args[0].checked_type.shape) == ifm_shape
+ assert op.args[0].checked_type.dtype == dtype
+ assert list(op.checked_type.shape) == ofm_shape
+ assert op.checked_type.dtype == dtype
+ assert op.attrs.pooling_type == pooling_type
+ assert list(op.attrs.strides) == strides_legalized
+ assert list(op.attrs.padding) == infra.compute_padding_shape(
+ ifm_shape, ofm_shape, padding, pool_shape, strides
+ )
+ assert list(op.attrs.padding) == infra.compute_padding_shape(
+ ifm_shape, ofm_shape, padding, pool_shape, strides_legalized
+ )
+ assert list(op.attrs.pool_shape) == pool_shape
+ assert op.attrs.ofm_channels == ifm_shape[3]
+ if activation_function == "RELU":
+ assert str(op.attrs.activation) == "CLIP"
+
+ if pooling_type == "MAX":
+ rewriter = legalize.MaxPoolingRewriter()
+ pattern_table = [
+ (
+ ethosu.MaxPool2DParams.composite_name,
+ ethosu.qnn_maxpool2d_pattern(),
+ lambda pat: ethosu.MaxPool2DParams(pat).is_valid(),
+ ),
+ ]
+
+ if pooling_type == "AVG":
+ rewriter = legalize.AvgPoolingRewriter()
+ pattern_table = [
+ (
+ ethosu.AvgPool2DParams.composite_name,
+ ethosu.qnn_avgpool2d_pattern(),
+ lambda pat: ethosu.AvgPool2DParams(pat).is_valid(),
+ ),
+ ]
+
+ tflite_graph = create_tflite_graph()
+ tflite_model = tflite.Model.Model.GetRootAsModel(tflite_graph, 0)
+
+ mod, _ = relay.frontend.from_tflite(
+ tflite_model,
+ shape_dict={"x": ifm_shape},
+ dtype_dict={"x": dtype},
+ )
+ mod = partition_ethosu_by_table(mod, pattern_table)
+
+ expected = expected_mod()
+ tvm.ir.assert_structural_equal(mod, expected)
+
+ mod["tvmgen_default_ethos_u_main_0"] = dataflow_pattern.rewrite(
+ rewriter, mod["tvmgen_default_ethos_u_main_0"]
+ )
+ verify(mod["tvmgen_default_ethos_u_main_0"])
+
+
@pytest.mark.parametrize("operator_type", ["ADD", "SUB", "MUL", "MIN", "MAX"])
@pytest.mark.parametrize(
"ifm_shape, ifm2_shape, reversed_operands",
