From 868eef2150039a5b23025719069d866914086bfd Mon Sep 17 00:00:00 2001
From: Digant Desai <digantdesai@meta.com>
Date: Fri, 25 Apr 2025 09:25:59 -0700
Subject: [PATCH] [ExecuTorch] Arm backend: Buckify cos test

Still no reference model

Differential Revision: [D73642290](https://our.internmc.facebook.com/intern/diff/D73642290/)

[ghstack-poisoned]
---
 backends/arm/test/ops/test_cos.py         |  7 ++++-
 backends/arm/test/targets.bzl             |  1 +
 backends/arm/test/tester/test_pipeline.py | 36 ++++++++++++++---------
 3 files changed, 29 insertions(+), 15 deletions(-)

diff --git a/backends/arm/test/ops/test_cos.py b/backends/arm/test/ops/test_cos.py
index 21902f51192..7cfd32d2bd2 100644
--- a/backends/arm/test/ops/test_cos.py
+++ b/backends/arm/test/ops/test_cos.py
@@ -6,8 +6,9 @@
 
 from typing import Tuple
 
-import torch
+import pytest
 
+import torch
 from executorch.backends.arm.test import common, conftest
 from executorch.backends.arm.test.tester.test_pipeline import (
     EthosU55PipelineBI,
@@ -37,24 +38,28 @@ def forward(self, x: torch.Tensor):
 
 
 @common.parametrize("test_data", test_data_suite)
+@pytest.mark.tosa_ref_model
 def test_cos_tosa_MI(test_data: Tuple):
     pipeline = TosaPipelineMI[input_t1](
         Cos(),
         (test_data,),
         aten_op,
         exir_op=[],
+        run_on_tosa_ref_model=conftest.is_option_enabled("tosa_ref_model"),
     )
     if conftest.get_option("tosa_version") == "1.0":
         pipeline.run()
 
 
 @common.parametrize("test_data", test_data_suite)
+@pytest.mark.tosa_ref_model
 def test_cos_tosa_BI(test_data: Tuple):
     pipeline = TosaPipelineBI[input_t1](
         Cos(),
         (test_data,),
         aten_op,
         exir_op=[],
+        run_on_tosa_ref_model=conftest.is_option_enabled("tosa_ref_model"),
     )
     pipeline.run()
 
diff --git a/backends/arm/test/targets.bzl b/backends/arm/test/targets.bzl
index 1993df0d091..832dcb3286c 100644
--- a/backends/arm/test/targets.bzl
+++ b/backends/arm/test/targets.bzl
@@ -17,6 +17,7 @@ def define_arm_tests():
         "ops/test_slice.py",
         "ops/test_sigmoid.py",
         "ops/test_tanh.py",
+        "ops/test_cos.py",
     ]
 
     # Quantization
diff --git a/backends/arm/test/tester/test_pipeline.py b/backends/arm/test/tester/test_pipeline.py
index 95d9898c19d..38d82b739e1 100644
--- a/backends/arm/test/tester/test_pipeline.py
+++ b/backends/arm/test/tester/test_pipeline.py
@@ -258,6 +258,8 @@ class TosaPipelineBI(BasePipelineMaker, Generic[T]):
        exir_ops: Exir dialect ops expected to be found in the graph after to_edge.
        if not using use_edge_to_transform_and_lower.
 
+       run_on_tosa_ref_model: Set to true to test the tosa file on the TOSA reference model.
+
        tosa_version: A string for identifying the TOSA version, see common.get_tosa_compile_spec for
                      options.
        use_edge_to_transform_and_lower: Selects betweeen two possible ways of lowering the module.
@@ -270,6 +272,7 @@ def __init__(
         test_data: T,
         aten_op: str | List[str],
         exir_op: Optional[str | List[str]] = None,
+        run_on_tosa_ref_model: bool = True,
         tosa_version: str = "TOSA-0.80+BI",
         symmetric_io_quantization: bool = False,
         use_to_edge_transform_and_lower: bool = True,
@@ -324,13 +327,14 @@ def __init__(
             suffix="quant_nodes",
         )
 
-        self.add_stage(
-            self.tester.run_method_and_compare_outputs,
-            atol=atol,
-            rtol=rtol,
-            qtol=qtol,
-            inputs=self.test_data,
-        )
+        if run_on_tosa_ref_model:
+            self.add_stage(
+                self.tester.run_method_and_compare_outputs,
+                atol=atol,
+                rtol=rtol,
+                qtol=qtol,
+                inputs=self.test_data,
+            )
 
 
 class TosaPipelineMI(BasePipelineMaker, Generic[T]):
@@ -345,6 +349,8 @@ class TosaPipelineMI(BasePipelineMaker, Generic[T]):
        exir_ops: Exir dialect ops expected to be found in the graph after to_edge.
        if not using use_edge_to_transform_and_lower.
 
+       run_on_tosa_ref_model: Set to true to test the tosa file on the TOSA reference model.
+
        tosa_version: A string for identifying the TOSA version, see common.get_tosa_compile_spec for
                      options.
        use_edge_to_transform_and_lower: Selects betweeen two possible ways of lowering the module.
@@ -357,6 +363,7 @@ def __init__(
         test_data: T,
         aten_op: str | List[str],
         exir_op: Optional[str | List[str]] = None,
+        run_on_tosa_ref_model: bool = True,
         tosa_version: str = "TOSA-0.80+MI",
         use_to_edge_transform_and_lower: bool = True,
         custom_path: str = None,
@@ -385,13 +392,14 @@ def __init__(
             suffix="quant_nodes",
         )
 
-        self.add_stage(
-            self.tester.run_method_and_compare_outputs,
-            atol=atol,
-            rtol=rtol,
-            qtol=qtol,
-            inputs=self.test_data,
-        )
+        if run_on_tosa_ref_model:
+            self.add_stage(
+                self.tester.run_method_and_compare_outputs,
+                atol=atol,
+                rtol=rtol,
+                qtol=qtol,
+                inputs=self.test_data,
+            )
 
 
 class EthosU55PipelineBI(BasePipelineMaker, Generic[T]):