Arm backend: Add support for non-folded quantized BatchNorm (pytorch#21710)

* Support quantized batch norm when it follows an operator that it
cannot be folded into.

Change-Id: I55ef21492c521cddbd8282ad9bfe17f0e2c573d8

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils
@Sebastian-Larsson @robell @rascani

Signed-off-by: Tom Allsop <tom.allsop@arm.com>

Author: tom-arm

backends/arm/_passes/arm_pass_manager.py

@@ -745,6 +745,7 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
                     ConvertEluFamilyToEluPass(tfa_pass=True),
                     DecomposeGroupNormPass(tfa_pass=True),
                     DecomposeLayerNormPass(tfa_pass=True),
+                    DecomposeBatchNormNoStatsPass(tfa_pass=True),
                     DecomposeVarPass(tfa_pass=True),
                     DecomposeMeanDimPass(graph_module, self.tosa_spec, tfa_pass=True),
                     DecomposeAdaptiveAvgPool2dPass(tfa_pass=True),

backends/arm/_passes/decompose_batch_norm_no_stats.py

@@ -61,6 +61,14 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
             ):
                 continue
 
+            if self.is_tfa_pass and getattr(node.args[0], "target", None) in (
+                torch.ops.aten.conv1d.default,
+                torch.ops.aten.conv2d.default,
+                torch.ops.aten.conv3d.default,
+                torch.ops.aten.convolution.default,
+            ):
+                continue
+
             if node.target in (
                 torch.ops.aten.batch_norm.default,
                 torch.ops.aten.native_batch_norm.default,
@@ -87,8 +95,17 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
             bias = args[2] if len(args) > 2 else None
             running_mean = args[3]
             running_var = args[4]
-            if len(args) > 6:
-                eps = args[6]
+            eps_index = (
+                7
+                if node.target
+                in (
+                    torch.ops.aten.batch_norm.default,
+                    torch.ops.aten.native_batch_norm.default,
+                )
+                else 6
+            )
+            if len(args) > eps_index:
+                eps = args[eps_index]
 
             # Determine shapes
             val = meta.get("val")
@@ -103,13 +120,20 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
             weights_shape[channel_axis] = shape[channel_axis]
             num_features = shape[channel_axis]
 
-            # Ops to use
-            sub_op = exir_ops.edge.aten.sub.Tensor
-            view_op = exir_ops.edge.aten.view_copy.default
-            full_op = exir_ops.edge.aten.full.default
-            add_op = exir_ops.edge.aten.add.Tensor
-            rsqrt_op = exir_ops.edge.aten.rsqrt.default
-            mul_op = exir_ops.edge.aten.mul.Tensor
+            if self.is_tfa_pass:
+                sub_op = torch.ops.aten.sub.Tensor
+                view_op = torch.ops.aten.view_copy.default
+                full_op = torch.ops.aten.full.default
+                add_op = torch.ops.aten.add.Tensor
+                rsqrt_op = torch.ops.aten.rsqrt.default
+                mul_op = torch.ops.aten.mul.Tensor
+            else:
+                sub_op = exir_ops.edge.aten.sub.Tensor
+                view_op = exir_ops.edge.aten.view_copy.default
+                full_op = exir_ops.edge.aten.full.default
+                add_op = exir_ops.edge.aten.add.Tensor
+                rsqrt_op = exir_ops.edge.aten.rsqrt.default
+                mul_op = exir_ops.edge.aten.mul.Tensor
 
             # Begin decomposition
             with graph_module.graph.inserting_before(node):

backends/arm/quantizer/quantizer_support.py

@@ -155,6 +155,7 @@ def check_pattern(cls, pattern):
 ALL_QPARAM_OP_PATTERNS = (
     [(target,) for target in _one_to_one]
     + ACTIVATION_FUNCTION_PATTERNS
+    + combo_pattern(BATCH_NORM_OPS)
     + CONV_OP_PATTERNS
     + LINEAR_OP_PATTERNS
     + BINARY_OP_PATTERNS

backends/arm/test/ops/test_batch_norm.py

@@ -15,7 +15,6 @@
 from executorch.backends.arm.test.tester.test_pipeline import (
     EthosU55PipelineINT,
     EthosU85PipelineINT,
-    OpNotSupportedPipeline,
     TosaPipelineFP,
     TosaPipelineINT,
     VgfPipeline,
@@ -87,17 +86,14 @@ def test_native_batch_norm_legit_no_training_tosa_FP(test_data: Tuple):
     pipeline.run()
 
 
-# TODO(MLETORCH-100: Quantized stand-alone batch norms)
-def test_native_batch_norm_legit_no_training_tosa_INT_not_delegated():
+def test_native_batch_norm_legit_no_training_tosa_INT():
     test_data, model_params = test_data_suite["rand_1_3_254_254"]()
-    OpNotSupportedPipeline[input_t1](
+    pipeline = TosaPipelineINT[input_t1](
         BatchNorm2d(*model_params),
         (test_data,),
-        {
-            "executorch_exir_dialects_edge__ops_aten__native_batch_norm_legit_no_training_default": 1
-        },
-        quantize=True,
-    ).run()
+        aten_op=[],
+    )
+    pipeline.run()
 
 
 @common.parametrize("test_data", test_data_suite)
@@ -120,31 +116,24 @@ def test_native_batch_norm_legit_no_training_vgf_quant(test_data: Tuple):
     pass
 
 
-# TODO(MLETORCH-100: Quantized stand-alone batch norms)
-def test_native_batch_norm_legit_no_training_u55_INT_not_delegated():
+def test_native_batch_norm_legit_no_training_u55_INT():
     test_data, model_params = test_data_suite["rand_1_3_254_254"]()
-    OpNotSupportedPipeline[input_t1](
+    pipeline = EthosU55PipelineINT[input_t1](
         BatchNorm2d(*model_params),
         (test_data,),
-        {
-            "executorch_exir_dialects_edge__ops_aten__native_batch_norm_legit_no_training_default": 1
-        },
-        quantize=True,
-        u55_subset=True,
-    ).run()
+        aten_ops=[],
+    )
+    pipeline.run()
 
 
-# TODO(MLETORCH-100: Quantized stand-alone batch norms)
-def test_native_batch_norm_legit_no_training_u85_INT_not_delegated():
+def test_native_batch_norm_legit_no_training_u85_INT():
     test_data, model_params = test_data_suite["rand_1_3_254_254"]()
-    OpNotSupportedPipeline[input_t1](
+    pipeline = EthosU85PipelineINT[input_t1](
         BatchNorm2d(*model_params),
         (test_data,),
-        {
-            "executorch_exir_dialects_edge__ops_aten__native_batch_norm_legit_no_training_default": 1
-        },
-        quantize=True,
-    ).run()
+        aten_ops=[],
+    )
+    pipeline.run()
 
 
 class BatchNorm2dConv(torch.nn.Module):