Merge branch 'release/0.4' into cherry-pick-6051-by-pytorch_bot_bot_

Author: shoumikhin

.ci/scripts/test_llama.sh

@@ -171,7 +171,7 @@ else
 fi
 
 # Check dtype.
-EXPORTED_MODEL_NAME="llama2"
+EXPORTED_MODEL_NAME="tinyllama_${MODE}_${DTYPE}"
 if [[ "${DTYPE}" == "fp16" ]]; then
   EXPORTED_MODEL_NAME="${EXPORTED_MODEL_NAME}_h"
 elif [[ "${DTYPE}" == "fp32" ]]; then

.ci/scripts/test_model.sh

@@ -155,30 +155,24 @@ test_model_with_qnn() {
 
   if [[ "${MODEL_NAME}" == "dl3" ]]; then
     EXPORT_SCRIPT=deeplab_v3
-    EXPORTED_MODEL_NAME=dlv3_qnn.pte
   elif [[ "${MODEL_NAME}" == "mv3" ]]; then
     EXPORT_SCRIPT=mobilenet_v3
-    EXPORTED_MODEL_NAME=mv3_qnn.pte
   elif [[ "${MODEL_NAME}" == "mv2" ]]; then
     EXPORT_SCRIPT=mobilenet_v2
-    EXPORTED_MODEL_NAME=mv2_qnn.pte
   elif [[ "${MODEL_NAME}" == "ic4" ]]; then
     EXPORT_SCRIPT=inception_v4
-    EXPORTED_MODEL_NAME=ic4_qnn.pte
   elif [[ "${MODEL_NAME}" == "ic3" ]]; then
     EXPORT_SCRIPT=inception_v3
-    EXPORTED_MODEL_NAME=ic3_qnn.pte
   elif [[ "${MODEL_NAME}" == "vit" ]]; then
     EXPORT_SCRIPT=torchvision_vit
-    EXPORTED_MODEL_NAME=vit_qnn.pte
   fi
 
   # Use SM8450 for S22, SM8550 for S23, and SM8560 for S24
   # TODO(guangyang): Make QNN chipset matches the target device
   QNN_CHIPSET=SM8450
 
   "${PYTHON_EXECUTABLE}" -m examples.qualcomm.scripts.${EXPORT_SCRIPT} -b ${CMAKE_OUTPUT_DIR} -m ${QNN_CHIPSET} --compile_only
-  EXPORTED_MODEL=./${EXPORT_SCRIPT}/${EXPORTED_MODEL_NAME}
+  EXPORTED_MODEL=$(find "./${EXPORT_SCRIPT}" -type f -name "${MODEL_NAME}*.pte" -print -quit)
 }
 
 test_model_with_coreml() {
@@ -187,8 +181,20 @@ test_model_with_coreml() {
     exit 1
   fi
 
-  "${PYTHON_EXECUTABLE}" -m examples.apple.coreml.scripts.export --model_name="${MODEL_NAME}"
+  DTYPE=float16
+
+  "${PYTHON_EXECUTABLE}" -m examples.apple.coreml.scripts.export --model_name="${MODEL_NAME}" --compute_precision "${DTYPE}"
   EXPORTED_MODEL=$(find "." -type f -name "${MODEL_NAME}*.pte" -print -quit)
+  # TODO:
+  if [ -n "$EXPORTED_MODEL" ]; then
+    EXPORTED_MODEL_WITH_DTYPE="${EXPORTED_MODEL%.pte}_${DTYPE}.pte"
+    mv "$EXPORTED_MODEL" "$EXPORTED_MODEL_WITH_DTYPE"
+    EXPORTED_MODEL="$EXPORTED_MODEL_WITH_DTYPE"
+    echo "Renamed file path: $EXPORTED_MODEL"
+  else
+    echo "No .pte file found"
+    exit 1
+  fi
 }
 
 if [[ "${BACKEND}" == "portable" ]]; then

backends/arm/TARGETS

@@ -8,7 +8,7 @@ python_library(
     typing = True,
     deps = [
         ":arm_backend",
-        "//executorch/backends/arm/passes:passes",
+        "//executorch/backends/arm/_passes:passes",
         "//executorch/exir:lib",
     ],
 )
@@ -27,7 +27,7 @@ python_library(
         ":arm_vela",
         "//executorch/backends/arm/operators:lib",
         "//executorch/backends/arm/operators:node_visitor",
-        "//executorch/backends/arm/passes:passes",
+        "//executorch/backends/arm/_passes:passes",
     ],
 )
 

backends/arm/passes/TARGETS renamed to backends/arm/_passes/TARGETS

@@ -8,20 +8,20 @@
 # pyre-unsafe
 
 import torch
-from executorch.backends.arm.passes.annotate_channels_last_dim_order_pass import (
+from executorch.backends.arm._passes.annotate_channels_last_dim_order_pass import (
     AnnotateChannelsLastDimOrder,
 )
-from executorch.backends.arm.passes.convert_expand_copy_to_repeat import (
+from executorch.backends.arm._passes.convert_expand_copy_to_repeat import (
     ConvertExpandCopyToRepeatPass,
 )
-from executorch.backends.arm.passes.convert_split_to_slice import (
+from executorch.backends.arm._passes.convert_split_to_slice import (
     ConvertSplitToSlicePass,
 )
-from executorch.backends.arm.passes.meandim_to_averagepool_pass import (
+from executorch.backends.arm._passes.meandim_to_averagepool_pass import (
     ConvertMeanDimToAveragePool,
 )
-from executorch.backends.arm.passes.remove_clone_pass import RemoveClonePass
-from executorch.backends.arm.passes.size_adjust_conv2d_pass import SizeAdjustConv2DPass
+from executorch.backends.arm._passes.remove_clone_pass import RemoveClonePass
+from executorch.backends.arm._passes.size_adjust_conv2d_pass import SizeAdjustConv2DPass
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 from executorch.exir.pass_manager import PassManager
 

backends/arm/passes/annotate_channels_last_dim_order_pass.py renamed to backends/arm/_passes/annotate_channels_last_dim_order_pass.py

@@ -0,0 +1,66 @@
+# Copyright 2024 Arm Limited and/or its affiliates.
+# All rights reserved.
+#
+# 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 Optional
+
+import torch
+
+from executorch.exir.dialects._ops import ops as exir_ops
+from torch._ops import OpOverload
+
+
+def create_node(
+    graph: torch.fx.Graph,
+    op_target: OpOverload,
+    args: tuple = (),
+    kwargs: Optional[dict] = None,
+    quantize: bool = False,
+    q_params: Optional[tuple] = None,
+):
+    """
+    Adds a node to 'graph'. graph.inserting_before/after() should be used before the call to decide where to insert the node.
+    If quantize is true and q_params is not None, a q dq pair is inserted after the newly created node.
+    """
+
+    node = graph.create_node(
+        "call_function",
+        op_target,
+        args=args,
+        kwargs=kwargs or {},
+    )
+    if quantize and q_params:
+        return insert_q_dq_pair(graph, node, q_params)
+    return node
+
+
+def insert_q_dq_pair(
+    graph: torch.fx.Graph,
+    anchor: torch.fx.Node,
+    q_params: tuple,
+):
+    """
+    Inserts a q dq node pair after the node 'anchor'.
+    """
+
+    with graph.inserting_after(anchor):
+        q = create_node(
+            graph=graph,
+            op_target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
+            args=(),  # We add the argument last
+        )
+        q.meta = anchor.meta
+    with graph.inserting_after(q):
+        dq = create_node(
+            graph=graph,
+            op_target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
+            args=(q,) + q_params,
+        )
+        dq.meta = q.meta
+    anchor.replace_all_uses_with(dq)
+    # We add this last so the replace all uses above does not replace the quantized
+    # node's first use
+    q.args = (anchor,) + q_params
+    return dq

backends/arm/passes/arm_pass_manager.py renamed to backends/arm/_passes/arm_pass_manager.py

@@ -0,0 +1,35 @@
+# Copyright 2024 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, PassResult
+
+
+class CastInt64ToInt32Pass(ExportPass):
+    def __init__(self, exported_program: torch.export.ExportedProgram):
+        super(CastInt64ToInt32Pass, self).__init__()
+        self.exported_program = exported_program
+
+    def _to_int32(self, graph_module: torch.fx.GraphModule):
+        for node in graph_module.graph.nodes:
+            fake_tensor = node.meta["val"]
+            if isinstance(fake_tensor, torch._subclasses.fake_tensor.FakeTensor):
+                if node.meta["val"].dtype == torch.int64:
+                    node.meta["val"] = node.meta["val"].to(torch.int32)
+                    buffer_name = (
+                        self.exported_program.graph_signature.inputs_to_buffers[
+                            node.name
+                        ]
+                    )
+                    new_tensor = self.exported_program.state_dict[buffer_name].to(
+                        torch.int32
+                    )
+                    self.exported_program.state_dict[buffer_name] = new_tensor
+
+    def call(self, graph_module: torch.fx.GraphModule):
+        self._to_int32(graph_module)
+        graph_module.recompile()
+        graph_module = super().call(graph_module).graph_module
+        return PassResult(graph_module, True)