Merge remote-tracking branch 'origin/main' into use-executorch-core

Author: kirklandsign

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()

backends/vulkan/runtime/graph/ops/glsl/conv1d.glsl

@@ -86,6 +86,9 @@ void main() {
   const int in_l = out_l * stride - padding;
   VEC4_T sum = VEC4_T(0);
 
+  const int out_c_packed_index = out_c >> 2;
+  const int out_c_packed_lane = out_c & 0x3;
+
   for (int in_c = c_start; in_c < c_end; ++in_c) {
     // "k" tracks the kernel's index for our input-kernel computation.
     // It reads out-of-bound zeros, but trying to avoid them complicates
@@ -103,16 +106,16 @@ void main() {
     // It is possible to further reduce the memory footprint by swapping the
     // dimensions, using x extent for out_channel, and y for kernel.
     for (int k = 0; k < kernel_size; k++) {
-      const ivec3 w_lpos = ivec3(k, in_c % in_group_size, out_c / 4);
+      const ivec3 w_lpos = ivec3(k, in_c % in_group_size, out_c_packed_index);
       const VEC4_T weight_texel = load_texel_lpos(kernel_in, w_lpos, kernel_axis_map);
-      VEC4_T weight = VEC4_T(weight_texel[out_c % 4]);
+      VEC4_T weight = VEC4_T(weight_texel[out_c_packed_lane]);
 
       const ivec3 in_pos = lpos_to_pos(ivec3(in_l + k * dilation, in_c, N), in_axis_map);
       sum = fma(weight, load_texel(t_in, in_pos), sum);
     }
   }
 
-  const VEC4_T bias = load_texel_lpos(bias_in, ivec3(out_c, 0, 0), bias_axis_map);
+  const VEC4_T bias = load_texel_lpos(bias_in, ivec3(out_c_packed_index, 0, 0), bias_axis_map);
   const ivec3 out_lpos = ivec3(out_l, out_c, N);
-  write_texel_lpos(t_out, out_lpos, op(sum + bias.x, out_min, out_max), out_axis_map);
+  write_texel_lpos(t_out, out_lpos, op(sum + bias[out_c_packed_lane], out_min, out_max), out_axis_map);
 }

backends/vulkan/runtime/graph/ops/impl/Convolution.cpp

@@ -483,7 +483,7 @@ void add_conv1d_node(
       weight,
       /*transposed = */ false,
       /*storage_type = */ utils::kTexture3D,
-      /*memory_layout = */ utils::kChannelsPacked);
+      /*memory_layout = */ utils::kWidthPacked);
 
   float out_min_val = 0.0f;
   float out_max_val = 0.0f;

tools/cmake/cmake_deps.toml

@@ -197,7 +197,6 @@ filters = [
 ]
 deps = [
   "executorch_core",
-  "extension_flat_tensor_schema",
 ]
 
 [targets.extension_module]
@@ -236,6 +235,8 @@ deps = [
   "extension_data_loader",
   "extension_flat_tensor",
   "extension_module",
+  "extension_data_loader",
+  "extension_flat_tensor",
   "extension_runner_util",
   "extension_tensor",
 ]