Merge branch 'main' into jz/add-thinking-toggle

Author: shoumikhin

.ci/scripts/build-qnn-sdk.sh

@@ -33,6 +33,7 @@ set_up_aot() {
   cmake .. \
       -DCMAKE_INSTALL_PREFIX=$PWD \
       -DEXECUTORCH_BUILD_QNN=ON \
+      -DANDROID_NATIVE_API_LEVEL=30 \
       -DQNN_SDK_ROOT=${QNN_SDK_ROOT} \
       -DEXECUTORCH_BUILD_DEVTOOLS=ON \
       -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \

Package.swift

@@ -96,10 +96,8 @@ let package = Package(
         .copy("resources/add.pte")
       ],
       linkerSettings: [
-        .linkedLibrary("c++"),
         .unsafeFlags([
-          "-Xlinker", "-force_load",
-          "-Xlinker", "cmake-out/kernels_portable.xcframework/macos-arm64/libkernels_portable_macos.a",
+          "-Xlinker", "-all_load",
         ])
       ]
     )

backends/apple/mps/setup.md

@@ -76,12 +76,12 @@ cd executorch
 ## Run the mv3 generated model using the mps_executor_runner
 
 ```bash
-./cmake-out/examples/apple/mps/mps_executor_runner --model_path mv3_mps_bundled_fp16.pte --bundled_program
+./cmake-out/examples/apple/mps/mps_executor_runner --model_path mv3_mps_float16_bundled.pte --bundled_program
 ```
 
 - You should see the following results. Note that no output file will be generated in this example:
 ```
-I 00:00:00.003290 executorch:mps_executor_runner.mm:286] Model file mv3_mps_bundled_fp16.pte is loaded.
+I 00:00:00.003290 executorch:mps_executor_runner.mm:286] Model file mv3_mps_float16_bundled.pte is loaded.
 I 00:00:00.003306 executorch:mps_executor_runner.mm:292] Program methods: 1
 I 00:00:00.003308 executorch:mps_executor_runner.mm:294] Running method forward
 I 00:00:00.003311 executorch:mps_executor_runner.mm:349] Setting up non-const buffer 1, size 606112.
@@ -118,7 +118,7 @@ python3 -m examples.apple.mps.scripts.mps_example --model_name="mv3" --generate_
 ```
 2. Run your Program on the ExecuTorch runtime and generate an [ETDump](../../../docs/source/etdump.md).
 ```
-./cmake-out/examples/apple/mps/mps_executor_runner --model_path mv3_mps_bundled_fp16.pte --bundled_program --dump-outputs
+./cmake-out/examples/apple/mps/mps_executor_runner --model_path mv3_mps_float16_bundled.pte --bundled_program --dump-outputs
 ```
 3. Create an instance of the Inspector API by passing in the ETDump you have sourced from the runtime along with the optionally generated ETRecord from step 1.
 ```bash

backends/cadence/aot/compiler_utils.py

@@ -109,12 +109,12 @@ def get_cascaded_ops(
     return nodes
 
 
-# Capture the effect of transpose op on incoming dimension order
-def get_transposed_dims(node: torch.fx.Node, dims: List[int]) -> List[int]:
+def get_transposed_dims(
+    node: torch.fx.Node, dims: Optional[List[int]] = None
+) -> List[int]:
     """
-    Given a transpose node, and the incoming dimension ordering of the input
-    tensor to the transpose node, return the net effect of transpose op on the
-    dimension order.
+    Applies the transposition as given by node onto the dimensions given in input
+    e.g (1, 2) on [a, b, c, d] would return [a, c, b, d]
     """
     assert node.target == exir_ops.edge.aten.transpose_copy.int
     # Assert that the dims is not empty
@@ -127,28 +127,22 @@ def get_transposed_dims(node: torch.fx.Node, dims: List[int]) -> List[int]:
     assert isinstance(transpose_dims1, int)
     dim0 = transpose_dims0 if transpose_dims0 >= 0 else transpose_dims0 + dim_len
     dim1 = transpose_dims1 if transpose_dims1 >= 0 else transpose_dims1 + dim_len
-    # Perform transpose on dimmension ordering (dims)
-    dims[dim0], dims[dim1] = dims[dim1], dims[dim0]
-    return dims
+    new_dims = list(dims)
+    new_dims[dim0], new_dims[dim1] = dims[dim1], dims[dim0]
+    return new_dims
 
 
-# Capture the effect of permute op on incoming dimension order
-def get_permuted_dims(node: torch.fx.Node, dims: Optional[Sequence[int]]) -> List[int]:
+def get_permuted_dims(node: torch.fx.Node, dims: List[int]) -> List[int]:
     """
-    Given a permute node, and the incoming dimension ordering of the input
-    tensor to the permute node, return the net effect of permute op on the
-    dimension order.
+    Applies the permutation as given by node onto the dimensions given in input
+    e.g (1, 2, 0)  on [a, b, c] would return [b, c, a]
     """
     assert node.target == exir_ops.edge.aten.permute_copy.default
     # Permute each index of the dimension ordering (dims)
     # pyre-fixme[6]: This combined typecheck isn't supported yet.
     permute_dims: List[int] = list(node.args[1])
     assert all(isinstance(x, int) for x in permute_dims)
-    # If the dims is empty, we can simply return the permute order
-    if not dims:
-        return permute_dims
-    dims = [dims[x] for x in permute_dims]
-    return dims
+    return [dims[x] for x in permute_dims]
 
 
 # Return the tensor of buffer/parameter op

backends/cadence/aot/fuse_ops.py

@@ -14,7 +14,7 @@
 import operator
 from collections import deque
 from numbers import Number
-from typing import cast, Sequence
+from typing import Any, Callable, cast
 
 # Import these for the cadence function signatures.
 import executorch.backends.cadence.aot.ops_registrations  # noqa: F401
@@ -881,9 +881,10 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
 
 
 @register_cadence_pass(CadencePassAttribute(opt_level=1))
-class FuseTransposeOpPairsPass(FuseOpPairsAcrossBranchesPass):
+class FuseTransposeOrPermuteOpPairsPass(FuseOpPairsAcrossBranchesPass):
     """
-    Fuse transpose op pairs to a single view op.
+    Fuse transpose or permute op pairs to a single view op.
+    (transpose or permutation) -> (quant or dequant) -> (transpose or permutation)
     """
 
     # A list of ops that can be bypassed when looking for a
@@ -907,42 +908,17 @@ def can_fuse_for_chain(
         if not super().can_fuse_for_chain(producer, consumer, consumer_op_packets):
             return False
 
-        def get_dims(node: torch.fx.Node) -> tuple[int, int]:
-            def canonicalize(dim: int) -> int:
-                if dim < 0:
-                    dim += len(node.meta["val"].shape)
-                return dim
-
-            return tuple(canonicalize(cast(int, d)) for d in node.args[1:3])
-
-        def is_equivalent(
-            shape: Sequence[int],
-            transpose0: tuple[int, int],
-            transpose1: tuple[int, int],
-        ) -> bool:
-            def permute_order(
-                order: Sequence[int], dims: tuple[int, int]
-            ) -> Sequence[int]:
-                new_order = list(order)
-                new_order[dims[0]], new_order[dims[1]] = (
-                    new_order[dims[1]],
-                    new_order[dims[0]],
-                )
-                return new_order
-
-            order = permute_order(range(len(shape)), transpose0)
-            order = permute_order(order, transpose1)
-
-            non_unit_dims = [dim for dim in range(len(shape)) if shape[dim] != 1]
-            non_unit_dims_permuted = [dim for dim in order if shape[dim] != 1]
-
-            return non_unit_dims == non_unit_dims_permuted
-
-        return is_equivalent(
-            cast(torch.fx.Node, producer.args[0]).meta["val"].shape,
-            get_dims(producer),
-            get_dims(consumer),
-        )
+        # checking that permut2(permut1(identify)) == identity
+        input_shape = cast(torch.fx.Node, producer.args[0]).meta["val"].shape
+        ident_dims = list(range(len(input_shape)))
+        # this mapping helps to handle both transpose and permutations
+        f: dict[Any, Callable] = {
+            exir_ops.edge.aten.transpose_copy.int: get_transposed_dims,
+            exir_ops.edge.aten.permute_copy.default: get_permuted_dims,
+        }
+        in_dims = f[producer.target](producer, ident_dims)
+        out_dims = f[consumer.target](consumer, in_dims)
+        return out_dims == ident_dims
 
     def get_fused_node(
         self,
@@ -960,11 +936,17 @@ def get_fused_node(
         return view
 
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
-        # Remove any dequantize op that has only quantize ops as its users.
+        # Remove any transpose/permutation op pair that cancel each other.
         self.find_and_fuse(
             graph_module,
-            producer_op_packets={exir_ops.edge.aten.transpose_copy},
-            consumer_op_packets={exir_ops.edge.aten.transpose_copy},
+            producer_op_packets={
+                exir_ops.edge.aten.transpose_copy,
+                exir_ops.edge.aten.permute_copy,
+            },
+            consumer_op_packets={
+                exir_ops.edge.aten.transpose_copy,
+                exir_ops.edge.aten.permute_copy,
+            },
             bypass_ops=self.bypass_ops,
         )
         result = super().call(graph_module)
@@ -1028,5 +1010,5 @@ class CadenceFuseOpsInGraph:
         FuseQuantDequantToRequantizePass,
         FuseMulIntoDequantPass,
         FuseFullThenReshapePass,
-        FuseTransposeOpPairsPass,
+        FuseTransposeOrPermuteOpPairsPass,
     ]

backends/cadence/aot/passes.py

@@ -14,7 +14,7 @@
 from executorch.backends.cadence.aot.fuse_ops import (
     CadenceFuseOpsInGraph,
     FuseFullThenReshapePass,
-    FuseTransposeOpPairsPass,
+    FuseTransposeOrPermuteOpPairsPass,
 )
 from executorch.backends.cadence.aot.pass_utils import (
     CadencePassAttribute,
@@ -83,7 +83,7 @@ def get_passes_in_default_order() -> List[ExportPass]:
         CadenceSimplifyOpsInGraph.passes,
         FinalizePipeline,
         FuseFullThenReshapePass,
-        FuseTransposeOpPairsPass,
+        FuseTransposeOrPermuteOpPairsPass,
         RemoveNopSliceOrViewOpPass,
     ]
     return pytree.tree_flatten(passes)[0]

backends/cadence/aot/replace_ops.py

@@ -2263,9 +2263,9 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
 
 
 @register_cadence_pass(CadencePassAttribute(opt_level=1))
-class ReplacePowWithMullPass(ExportPass):
+class ReplacePowWithMulPass(ExportPass):
     """
-    Replace the pow op with degree 2 for a mul op.
+    Replace the pow op for a mul op.
     """
 
     def call_operator(
@@ -2275,19 +2275,32 @@ def call_operator(
         kwargs: Dict[str, Argument],
         meta: NodeMetadata,
     ) -> ProxyValue:
-        # TODO(eigen): Add support for other degrees.
-        if (
-            op
-            not in {
-                exir_ops.edge.aten.pow.Scalar,
+        if not (
+            len(args) > 1
+            and isinstance(args[1], int)
+            and cast(int, args[1]) > 1
+            and cast(int, args[1]) < 5
+            and op
+            in {
+                exir_ops.edge.aten.pow.Tensor_Scalar,
             }
-            or args[0] != 2
         ):
             return super().call_operator(op, args, kwargs, meta)
 
+        x = args[0]
+        exponent = cast(int, args[1])
+
+        if exponent > 2:
+            for _ in range(exponent, 2, -1):
+                x = super().call_operator(
+                    exir_ops.edge.aten.mul.Tensor,
+                    (x, args[0]),
+                    {},
+                    meta,
+                )
         return super().call_operator(
             exir_ops.edge.aten.mul.Tensor,
-            (args[1], args[1]),
+            (x, args[0]),
             {},
             meta,
         )
@@ -2429,5 +2442,5 @@ class CadenceReplaceOpsInGraph:
         ReplaceWhereWithFullArgsWithWhereScalar,
         ReplaceGeluWithApproximateGeluPass,
         ReplaceSplitWithSlicePass,
-        ReplacePowWithMullPass,
+        ReplacePowWithMulPass,
     ]