Cortex_m backend: Loosen edge op check. (#13550)

The pass checked that all ops were edge ops to detect if the pass was
ran before lowering to edge. However, there are cases where aten ops
survive after edge lowering, notably torch.ops.tensor_scalar. This
shouldn't crash the pass. Instead, only check that q/dq ops are edge
ops.

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218

Signed-off-by: Erik Lundell <[email protected]>

Author: Erik-Lundell

backends/cortex_m/passes/replace_quant_nodes_pass.py

@@ -1,5 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2025 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -10,7 +11,6 @@
 import torch
 
 from executorch.exir.dialects._ops import ops as exir_ops
-from executorch.exir.dialects.edge._ops import EdgeOpOverload
 from executorch.exir.pass_base import ExportPass, NodeMetadata, ProxyValue
 
 
@@ -40,6 +40,10 @@ def __init__(self):
                 "qualifier": self._is_qualified_int8_node,
             },
         }
+        self.disallowed_targets = {
+            torch.ops.quantized_decomposed.dequantize_per_tensor.default,
+            torch.ops.quantized_decomposed.quantize_per_tensor.default,
+        }
 
     def call_operator(
         self,
@@ -48,9 +52,10 @@ def call_operator(
         kwargs: Dict[str, object],
         meta: NodeMetadata,
     ) -> ProxyValue:
-        assert isinstance(
-            op, EdgeOpOverload
-        ), "Op must be an EdgeOpOverload. Run this pass after to_edge()."
+        if op in self.disallowed_targets:
+            raise RuntimeError(
+                f"Found unexpected aten op '{op}'. Make sure you run this pass after lowering to edge."
+            )
 
         if op in self.op_replacements and self.op_replacements[op]["qualifier"](args):
             return super().call_operator(

backends/cortex_m/test/test_replace_quant_nodes.py

@@ -9,14 +9,16 @@
 from dataclasses import dataclass
 from typing import Optional
 
-import executorch
 import executorch.backends.cortex_m.ops.operators  # noqa
 
+import executorch.exir
+
 import torch
 from executorch.backends.cortex_m.passes.replace_quant_nodes_pass import (
     ReplaceQuantNodesPass,
 )
 from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.program._program import _transform
 from torch.export import export
 from torch.fx import GraphModule
 from torchao.quantization.pt2e.observer import HistogramObserver
@@ -128,11 +130,18 @@ def forward(self, x):
         # Step 1: Export and quantize the model
         exported_model = export(model.eval(), example_inputs, strict=True).module()
         prepared_model = prepare_pt2e(exported_model, AddQuantizer())
+        prepared_model(*example_inputs)
         quantized_model = convert_pt2e(prepared_model)
 
         # Step 2: Export to EXIR
         exported = export(quantized_model, example_inputs, strict=True)
 
+        # The pass should raise an Exception if ran before to_edge.
+        with self.assertRaisesRegex(
+            Exception, "An error occurred when running the 'ReplaceQuantNodesPass' pass"
+        ):
+            _transform(exported, ReplaceQuantNodesPass())
+
         # Step 3: Convert to Edge
         edge_program = executorch.exir.to_edge(
             exported,