Update base for Update on "Xnnpack test for program-data separation"

Add xnnpack test for program-data separation

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

[ghstack-poisoned]

Author: lucylq

backends/arm/_passes/__init__.py

@@ -19,6 +19,7 @@
 from .convert_squeezes_to_view import ConvertSqueezesToViewPass  # noqa
 from .convert_to_clamp import ConvertToClampPass  # noqa
 from .decompose_batchnorm_pass import DecomposeBatchNormPass  # noqa
+from .decompose_cosine_similarity_pass import DecomposeCosineSimilarityPass  # noqa
 from .decompose_div_pass import DecomposeDivPass  # noqa
 from .decompose_gelu_pass import DecomposeGeluPass  # noqa
 from .decompose_layernorm_pass import DecomposeLayerNormPass  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -24,6 +24,7 @@
     ConvertSqueezesToViewPass,
     ConvertToClampPass,
     DecomposeBatchNormPass,
+    DecomposeCosineSimilarityPass,
     DecomposeDivPass,
     DecomposeGeluPass,
     DecomposeLayerNormPass,
@@ -205,6 +206,7 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
         self.add_pass(DecomposeVarPass())
         self.add_pass(DecomposeMeanDimPass())
         self.add_pass(DecomposeNotEqualPass())
+        self.add_pass(DecomposeCosineSimilarityPass())
         self.add_pass(DecomposeDivPass())
         self.add_pass(DecomposeLeakyReLUPass())
         self.add_pass(DecomposeSqrtPass())

backends/arm/_passes/decompose_cosine_similarity_pass.py

@@ -0,0 +1,75 @@
+# 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.
+
+import torch
+from executorch.exir.pass_base import ExportPass
+
+torch_cosine_similarity = (torch.ops.aten.cosine_similarity.default,)
+
+
+class DecomposeCosineSimilarityPass(ExportPass):
+    """
+    Decomposition of aten.cosine_similarity:
+
+      dot    = sum(mul(x1, x2), dims, keepdim=False)
+      norm   = pow( sum(mul(x, x), dims, keepdim=False), 0.5 )
+      eps    = full( (), eps_scalar )
+      n1c    = max(norm1, eps)
+      n2c    = max(norm2, eps)
+      denom  = mul(n1c, n2c)
+      out    = div(dot, denom)
+    """
+
+    def call_operator(self, op, args, kwargs, meta):
+        if op not in torch_cosine_similarity:
+            return super().call_operator(op, args, kwargs, meta)
+
+        x1, x2 = args[0], args[1]
+        dim = kwargs.get("dim", 1)
+        eps = kwargs.get("eps", 1e-8)
+        dims = [dim] if isinstance(dim, int) else list(dim)
+
+        # 1) dot
+        prod = super().call_operator(torch.ops.aten.mul.Tensor, (x1, x2), {}, meta)
+        dot = super().call_operator(
+            torch.ops.aten.sum.dim_IntList, (prod, dims, False), {}, meta
+        )
+
+        # 2a) norm1 = pow(sum(x1*x1), 0.5)
+        x1_sq = super().call_operator(torch.ops.aten.mul.Tensor, (x1, x1), {}, meta)
+        s1 = super().call_operator(
+            torch.ops.aten.sum.dim_IntList, (x1_sq, dims, False), {}, meta
+        )
+        norm1 = super().call_operator(
+            torch.ops.aten.pow.Tensor_Scalar, (s1, 0.5), {}, meta
+        )
+
+        # 2b) norm2 = pow(sum(x2*x2), 0.5)
+        x2_sq = super().call_operator(torch.ops.aten.mul.Tensor, (x2, x2), {}, meta)
+        s2 = super().call_operator(
+            torch.ops.aten.sum.dim_IntList, (x2_sq, dims, False), {}, meta
+        )
+        norm2 = super().call_operator(
+            torch.ops.aten.pow.Tensor_Scalar, (s2, 0.5), {}, meta
+        )
+
+        # 3) eps scalar - we need to broadcast ourselves as TOSA dont do this for scalar
+        eps_t = super().call_operator(
+            torch.ops.aten.full_like.default, (norm1, eps), {}, meta
+        )
+
+        # 4) clamp to avoid zero division
+        n1c = super().call_operator(
+            torch.ops.aten.maximum.default, (norm1, eps_t), {}, meta
+        )
+        n2c = super().call_operator(
+            torch.ops.aten.maximum.default, (norm2, eps_t), {}, meta
+        )
+
+        # 5) denom and divide
+        denom = super().call_operator(torch.ops.aten.mul.Tensor, (n1c, n2c), {}, meta)
+        out = super().call_operator(torch.ops.aten.div.Tensor, (dot, denom), {}, meta)
+
+        return out

backends/arm/test/models/test_nn_functional.py

@@ -106,7 +106,6 @@ def test_nn_functional_MI(test_data):
 
 x_fails = {
     "normalize": "MLETORCH-852: Support aten.index_put.default",
-    "cosine_similarity": "MLETORCH-854: Support aten.linalg_vector_norm.default",
     "unfold": "Int64 input && MLETORCH-827: Support aten.index.Tensor",
     "fold": "Int64 input && MLETORCH-827: Support aten.index_put.default",
 }

backends/arm/test/passes/test_decompose_cosine_similarity_pass.py

@@ -0,0 +1,52 @@
+# 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.
+
+from typing import Tuple
+
+import torch
+
+from executorch.backends.arm._passes.decompose_cosine_similarity_pass import (
+    DecomposeCosineSimilarityPass,
+)
+from executorch.backends.arm.test import common
+from executorch.backends.arm.test.tester.test_pipeline import PassPipeline
+
+input_t = Tuple[torch.Tensor, torch.Tensor]
+
+
+class CosineSimilarityModel(torch.nn.Module):
+    def get_inputs(self) -> input_t:
+        return (torch.rand(2, 3, 4), torch.rand(2, 3, 4))
+
+    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
+        return torch.cosine_similarity(x1, x2, dim=1, eps=1e-6)
+
+
+modules = {"cosine_basic": CosineSimilarityModel()}
+
+
+@common.parametrize("module", modules)
+def test_decompose_cosine_similarity_tosa_BI(module):
+
+    ops_after_pass = {
+        "executorch_exir_dialects_edge__ops_aten_mul_Tensor": 5,
+        "executorch_exir_dialects_edge__ops_aten_sum_dim_IntList": 3,
+        "executorch_exir_dialects_edge__ops_aten_pow_Tensor_Scalar": 2,
+        "executorch_exir_dialects_edge__ops_aten_full_like_default": 1,
+        "executorch_exir_dialects_edge__ops_aten_maximum_default": 2,
+        "executorch_exir_dialects_edge__ops_aten_reciprocal_default": 1,
+    }
+
+    pipeline = PassPipeline[input_t](
+        module,
+        module.get_inputs(),
+        tosa_version="TOSA-0.80+BI",
+        ops_before_pass=None,
+        ops_not_before_pass=None,
+        ops_after_pass=ops_after_pass,
+        ops_not_after_pass=None,
+        pass_list=[DecomposeCosineSimilarityPass],
+    )
+    pipeline.run()

examples/models/llama/runner/CMakeLists.txt

@@ -53,7 +53,7 @@ else()
 endif()
 
 set(llama_runner_deps executorch_core extension_data_loader extension_module
-                      extension_tensor
+                      extension_tensor extension_flat_tensor
 )
 
 target_link_libraries(llama_runner PUBLIC ${llama_runner_deps})

examples/models/llava/runner/CMakeLists.txt

@@ -41,7 +41,7 @@ add_subdirectory(
 add_library(llava_runner STATIC ${_llava_runner__srcs})
 
 set(llava_runner_deps executorch_core extension_data_loader extension_llm_runner
-                      extension_module extension_tensor
+                      extension_module extension_tensor extension_flat_tensor
 )
 
 target_link_libraries(llava_runner PUBLIC ${llava_runner_deps})

extension/android/CMakeLists.txt

@@ -72,6 +72,7 @@ list(
   link_libraries
   executorch
   extension_data_loader
+  extension_flat_tensor
   extension_module
   extension_runner_util
   extension_tensor

extension/llm/runner/CMakeLists.txt

@@ -43,7 +43,7 @@ target_include_directories(
 
 add_library(extension_llm_runner STATIC ${_extension_llm_runner__srcs})
 
-set(runner_deps executorch extension_module extension_tensor)
+set(runner_deps executorch_core extension_module extension_tensor)
 
 target_link_libraries(extension_llm_runner PUBLIC ${runner_deps})
 

tools/cmake/cmake_deps.toml

@@ -196,9 +196,7 @@ filters = [
   ".cpp$",
 ]
 deps = [
-  "extension_flat_tensor_schema",
   "executorch_core",
-  "executorch",
 ]
 
 [targets.extension_module]
@@ -209,9 +207,9 @@ filters = [
   ".cpp$",
 ]
 deps = [
-  "executorch",
   "executorch_core",
   "extension_data_loader",
+  "extension_flat_tensor",
 ]
 
 [targets.extension_runner_util]
@@ -233,9 +231,12 @@ filters = [
   ".cpp$",
 ]
 deps = [
-  "executorch",
   "executorch_core",
+  "extension_data_loader",
+  "extension_flat_tensor",
   "extension_module",
+  "extension_data_loader",
+  "extension_flat_tensor",
   "extension_runner_util",
   "extension_tensor",
 ]
@@ -248,7 +249,6 @@ filters = [
   ".cpp$",
 ]
 deps = [
-  "executorch",
   "executorch_core",
 ]
 
@@ -260,7 +260,6 @@ filters = [
   ".cpp$",
 ]
 deps = [
-  "executorch",
   "executorch_core",
 ]
 
@@ -452,7 +451,9 @@ deps = [
   "executorch",
   "executorch_core",
   "extension_data_loader",
+  "extension_flat_tensor",
   "extension_module",
+  "extension_tensor",
   "extension_threadpool",
   "optimized_cpublas",
   "portable_kernels",

tools/cmake/executorch-config.cmake

@@ -66,6 +66,7 @@ set(lib_list
     etdump
     bundled_program
     extension_data_loader
+    extension_flat_tensor
     ${FLATCCRT_LIB}
     coreml_util
     coreml_inmemoryfs