Arm backend: Enable disallowing transforming prior to quantization (pytorch#15958)

ReplaceInfValues is provided with extra logic to handle this new flag. A
file test_replace_inf_values_pass.py that test the pass in general as
well as the flag has been added.

A test that verifies the new capability of the quantizer has been added
to a new file test_partial_quantization.py.

Signed-off-by: Martin Lindström <Martin.Lindstroem@arm.com>
Co-authored-by: Martin Lindström <Martin.Lindstroem@arm.com>

Author: martinlsm

backends/arm/_passes/arm_pass.py

@@ -8,6 +8,7 @@
 from abc import abstractmethod
 from typing import Any, List, Optional, Set, Type
 
+from executorch.backends.arm.constants import DISALLOW_TFA_META_KEY
 from executorch.exir.pass_base import ExportPass, NodeMetadata
 from torch.fx import GraphModule
 from torch.fx.passes.infra.pass_base import PassResult
@@ -16,9 +17,23 @@
 class ArmPass(ExportPass):
     """Base class for Arm passes"""
 
-    def __init__(self) -> None:
-        super().__init__()
+    def __init__(self, tfa_pass: bool = False, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
         self.submodule_depth = 0
+        self.is_tfa_pass = tfa_pass
+
+    def allowed_to_transform(self, meta: NodeMetadata | dict[str, Any]) -> bool:
+        if not self.is_tfa_pass:
+            return True
+
+        if isinstance(meta, NodeMetadata):
+            meta_dict = meta.data
+        else:
+            meta_dict = meta
+
+        disallow_tfa = meta_dict.get(DISALLOW_TFA_META_KEY, False)
+
+        return not disallow_tfa
 
     @property
     @abstractmethod

backends/arm/_passes/arm_pass_manager.py

@@ -376,65 +376,65 @@ def transform_to_backend_pipeline(
 
     def transform_for_annotation_pipeline(self, graph_module: GraphModule):
         # Preprocessing passes
-        self.add_pass(RemoveGraphAssertsPass())
+        self.add_pass(RemoveGraphAssertsPass(tfa_pass=True))
 
         # Transformation passes (pre scalar -> tensor)
         self.add_passes(
             [
-                DecomposeSelectScatterPass(),
-                ConvertInt64ConstOpsToInt32Pass(),
-                ConvertInt64OutputOpsToInt32Pass(),
-                InsertInt32CastsAfterInt64PlaceholdersPass(),
-                DecomposeEmbeddingPass(),
-                DecomposeScaledDotProductAttentionPass(),
-                DecomposeRoundPass(),
-                DecomposeLogitPass(),
-                PromoteBoolOperandsPass(),
-                DecomposeSignPass(),
-                DecomposeAddmmPass(),
-                DecomposeRemainderPass(),
-                DecomposeFloorDividePass(),
-                DecomposeDivTensorModePass(),
+                DecomposeSelectScatterPass(tfa_pass=True),
+                ConvertInt64ConstOpsToInt32Pass(tfa_pass=True),
+                ConvertInt64OutputOpsToInt32Pass(tfa_pass=True),
+                InsertInt32CastsAfterInt64PlaceholdersPass(tfa_pass=True),
+                DecomposeEmbeddingPass(tfa_pass=True),
+                DecomposeScaledDotProductAttentionPass(tfa_pass=True),
+                DecomposeRoundPass(tfa_pass=True),
+                DecomposeLogitPass(tfa_pass=True),
+                PromoteBoolOperandsPass(tfa_pass=True),
+                DecomposeSignPass(tfa_pass=True),
+                DecomposeAddmmPass(tfa_pass=True),
+                DecomposeRemainderPass(tfa_pass=True),
+                DecomposeFloorDividePass(tfa_pass=True),
+                DecomposeDivTensorModePass(tfa_pass=True),
             ]
         )
 
         # Scalars -> tensors
         self.add_passes(
             [
-                ReplaceScalarWithTensorByProfilePass(),
-                ScalarsToAttributePass(),
+                ReplaceScalarWithTensorByProfilePass(tfa_pass=True),
+                ScalarsToAttributePass(tfa_pass=True),
             ]
         )
 
         # Transformation passes (post scalar removal)
         self.add_passes(
             [
-                NormalizeWhileInitialArgsPass(use_exir_clone=False),
-                DecomposeAddSubAlphaPass(),
-                DecomposeGroupNormPass(),
-                DecomposeLayerNormPass(),
-                DecomposeVarPass(),
-                DecomposeMeanDimPass(graph_module, self.tosa_spec),
-                DecomposeNotEqualPass(),
-                DecomposeCosineSimilarityPass(),
-                DecomposeGluPass(),
-                DecomposeDivPass(),
-                DecomposeLeakyReLUPass(),
-                DecomposeLinalgVectorNormPass(),
-                DecomposeSqrtPass(),
-                DecomposeSiluPass(),
-                DecomposeAvgPool2dPass(),
-                DecomposeSoftmaxUnstablePass(),
-                DecomposeSoftmaxPass(),
-                ConvertMinMaxPass(),
+                NormalizeWhileInitialArgsPass(use_exir_clone=False, tfa_pass=True),
+                DecomposeAddSubAlphaPass(tfa_pass=True),
+                DecomposeGroupNormPass(tfa_pass=True),
+                DecomposeLayerNormPass(tfa_pass=True),
+                DecomposeVarPass(tfa_pass=True),
+                DecomposeMeanDimPass(graph_module, self.tosa_spec, tfa_pass=True),
+                DecomposeNotEqualPass(tfa_pass=True),
+                DecomposeCosineSimilarityPass(tfa_pass=True),
+                DecomposeGluPass(tfa_pass=True),
+                DecomposeDivPass(tfa_pass=True),
+                DecomposeLeakyReLUPass(tfa_pass=True),
+                DecomposeLinalgVectorNormPass(tfa_pass=True),
+                DecomposeSqrtPass(tfa_pass=True),
+                DecomposeSiluPass(tfa_pass=True),
+                DecomposeAvgPool2dPass(tfa_pass=True),
+                DecomposeSoftmaxUnstablePass(tfa_pass=True),
+                DecomposeSoftmaxPass(tfa_pass=True),
+                ConvertMinMaxPass(tfa_pass=True),
             ]
         )
 
         # Postprocessing passes
         self.add_passes(
             [
-                ReplaceInfAndLimitValuesPass(),
-                DecomposeMaskedFillPass(),
+                ReplaceInfAndLimitValuesPass(tfa_pass=True),
+                DecomposeMaskedFillPass(tfa_pass=True),
             ]
         )
 

backends/arm/_passes/cast_int64_pass.py

@@ -23,8 +23,8 @@ class CastInt64BuffersToInt32Pass(ArmPass):
 
     _passes_required_after: Set[Type[ExportPass]] = set()
 
-    def __init__(self, exported_program: ExportedProgram):
-        super().__init__()
+    def __init__(self, exported_program: ExportedProgram, *args, **kwargs):
+        super().__init__(*args, **kwargs)
         self.exported_program = exported_program
 
     def _assert_within_int32(self, tensor: torch.Tensor, node: torch.fx.Node):

backends/arm/_passes/convert_elu_params.py

@@ -35,7 +35,7 @@ def call(self, graph_module: torch.fx.GraphModule):
             is_quantized = (
                 input_node.op == "call_function" and input_node.target in DQ_OPS
             )
-            if not is_quantized:
+            if not is_quantized or not self.allowed_to_transform(node.meta):
                 continue
             with graph.inserting_after(node):
                 replace_node = create_node(

backends/arm/_passes/convert_minmax_pass.py

@@ -41,6 +41,15 @@ class ConvertMinMaxPass(ArmPass):
 
     _passes_required_after: Set[Type[ExportPass]] = {ConvertSqueezesToViewPass}
 
+    _TARGET_OPS = {
+        exir_ops.edge.aten.amax.default,
+        exir_ops.edge.aten.amin.default,
+        exir_ops.edge.aten.max.dim,
+        exir_ops.edge.aten.min.dim,
+        torch.ops.aten.max.dim,
+        torch.ops.aten.min.dim,
+    }
+
     def check_argmax(self, node):
         """
         Raises a RuntimeError if the argmax value returned by the min/max op is used in the graph.
@@ -88,16 +97,11 @@ def get_variables(self, node):
     def call(self, graph_module: torch.fx.GraphModule):
         modified = False
         for node in graph_module.graph.nodes:
-            if node.op != "call_function":
-                continue
-            if node.target not in [
-                exir_ops.edge.aten.amax.default,
-                exir_ops.edge.aten.amin.default,
-                exir_ops.edge.aten.max.dim,
-                exir_ops.edge.aten.min.dim,
-                torch.ops.aten.max.dim,
-                torch.ops.aten.min.dim,
-            ]:
+            if (
+                node.op != "call_function"
+                or node.target not in ConvertMinMaxPass._TARGET_OPS
+                or not self.allowed_to_transform(node.meta)
+            ):
                 continue
 
             self.check_argmax(

backends/arm/_passes/convert_to_clamp_pass.py

@@ -33,7 +33,7 @@ class ConvertToClampPass(ArmPass):
     _passes_required_after: Set[Type[ExportPass]] = {QuantizeClampArgumentsPass}
 
     def call_operator(self, op, args, kwargs, meta):
-        if op not in edge_operators:
+        if op not in edge_operators or not self.allowed_to_transform(meta):
             return super().call_operator(op, args, kwargs, meta)
 
         return super().call_operator(

backends/arm/_passes/decompose_addmm_pass.py

@@ -49,7 +49,7 @@ class DecomposeAddmmPass(ArmPass):
     }
 
     def call_operator(self, op, args, kwargs, meta):
-        if op not in [edge_addmm, aten_addmm]:
+        if op not in [edge_addmm, aten_addmm] or not self.allowed_to_transform(meta):
             return super().call_operator(op, args, kwargs, meta)
 
         input, mat1, mat2 = args

backends/arm/_passes/decompose_avg_pool2d_pass.py

@@ -43,7 +43,9 @@ class DecomposeAvgPool2dPass(ArmPass):
     _passes_required_after: Set[Type[ExportPass]] = {ComputeConstantOpsAOTPass}
 
     def call_operator(self, op, args, kwargs, meta):
-        if op not in (edge_div_ops + aten_div_ops):
+        if op not in (edge_div_ops + aten_div_ops) or not self.allowed_to_transform(
+            meta
+        ):
             return super().call_operator(op, args, kwargs, meta)
 
         full_op, cat_op, avgpool_op, mul_op = get_decomposition(op)

backends/arm/_passes/decompose_batch_norm_no_stats.py

@@ -53,7 +53,11 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
         )
 
         for node in graph_module.graph.nodes:
-            if node.op != "call_function" or node.target not in bn_ops:
+            if (
+                node.op != "call_function"
+                or node.target not in bn_ops
+                or not self.allowed_to_transform(node.meta)
+            ):
                 continue
 
             if node.target in (

backends/arm/_passes/decompose_cosine_similarity_pass.py

@@ -40,7 +40,7 @@ class DecomposeCosineSimilarityPass(ArmPass):
     }
 
     def call_operator(self, op, args, kwargs, meta):
-        if op not in torch_cosine_similarity:
+        if op not in torch_cosine_similarity or not self.allowed_to_transform(meta):
             return super().call_operator(op, args, kwargs, meta)
 
         x1, x2 = args[0], args[1]