Arm backend: Minor improvements in Arm testing. (#8337)

* Check that reference model was run in TosaReferenceModelDispatch

If the if statment looking for executorch_call_delegates fails,
the delegate can run the model in eager mode and still produce correct results
without the user noticing. We need to make sure we actually run
the reference model.

* Add markers to flaky tests

Author: Erik-Lundell

backends/arm/test/ops/test_bmm.py

@@ -124,6 +124,7 @@ def test_bmm_tosa_MI(self, test_data_generator: Callable[[], Tuple]):
         self._test_bmm_tosa_MI_pipeline(self.BMM(), test_data)
 
     @parameterized.expand(BMMSingleInput.test_data_generators)
+    @pytest.mark.flaky  # TODO: Investigate flakyness (MLETORCH-534)
     def test_bmm_single_input_tosa_MI(self, test_data_generator: Callable[[], Tuple]):
         test_data = test_data_generator()
         self._test_bmm_tosa_MI_pipeline(self.BMMSingleInput(), test_data)
@@ -144,6 +145,7 @@ def test_bmm_tosa_BI(self, test_data_generator: Callable[[], Tuple]):
         self._test_bmm_tosa_BI_pipeline(self.BMM(), test_data)
 
     @parameterized.expand(BMMSingleInput.test_data_generators)
+    @pytest.mark.flaky  # TODO: Investigate flakyness (MLETORCH-534)
     def test_bmm_single_input_tosa_BI(self, test_data_generator: Callable[[], Tuple]):
         test_data = test_data_generator()
         self._test_bmm_tosa_BI_pipeline(self.BMMSingleInput(), test_data)

backends/arm/test/ops/test_mm.py

@@ -115,6 +115,7 @@ def test_mm_tosa_MI(self, test_data_generator: Callable[[], Tuple]):
         self._test_mm_tosa_MI_pipeline(self.MM(), test_data)
 
     @parameterized.expand(MMSingleInput.test_data_generators)
+    @pytest.mark.flaky  # TODO: Investigate flakyness (MLETORCH-534)
     def test_mm_single_input_tosa_MI(self, test_data_generator: Callable[[], Tuple]):
         test_data = test_data_generator()
         self._test_mm_tosa_MI_pipeline(self.MMSingleInput(), test_data)
@@ -126,6 +127,7 @@ def test_mm_tosa_BI(self, test_data_generator: Callable[[], Tuple]):
         self._test_mm_tosa_BI_pipeline(self.MM(), test_data)
 
     @parameterized.expand(MMSingleInput.test_data_generators)
+    @pytest.mark.flaky  # TODO: Investigate flakyness (MLETORCH-534)
     def test_mm_single_input_tosa_BI(self, test_data_generator: Callable[[], Tuple]):
         test_data = test_data_generator()
         self._test_mm_tosa_BI_pipeline(self.MMSingleInput(), test_data)

backends/arm/test/runner_utils.py

@@ -157,6 +157,10 @@ def get_output_quantization_params(
 class TosaReferenceModelDispatch(TorchFunctionMode):
     """A context manager for executing call_delegate nodes using the reference model"""
 
+    def __init__(self):
+        self.ran_tosa_dispatch = False
+        super().__init__()
+
     def _tosa_dispatch(self, lowered_backend_module: LoweredBackendModule, inputs):
         tosa_buffer = lowered_backend_module.processed_bytes
         compile_specs = lowered_backend_module.compile_specs
@@ -168,13 +172,21 @@ def _tosa_dispatch(self, lowered_backend_module: LoweredBackendModule, inputs):
 
         return run_tosa_graph(tosa_buffer, tosa_version, inputs)
 
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        super().__exit__(exc_type, exc_val, exc_tb)
+        if not self.ran_tosa_dispatch:
+            raise RuntimeError(
+                "Ran model with TosaReferenceModelDispatch but never ran ArmBackend delegate."
+            )
+
     def __torch_function__(self, func, types, args=..., kwargs=None):
         if func is torch._higher_order_ops.executorch_call_delegate:
             lowered_backend_module = cast(LoweredBackendModule, args[0])
             if lowered_backend_module.backend_id == "ArmBackend":
+                self.ran_tosa_dispatch = True
                 return self._tosa_dispatch(lowered_backend_module, args[1:])
             else:
-                logger.warning(
+                raise RuntimeError(
                     f"Ran model with TosaReferenceModelDispatch but call_delegate with {lowered_backend_module.backend_id=} != 'ArmBackend'."
                 )