Merge branch 'main' into marlin-aot-arm-compiler

Author: martinlsm

.ci/scripts/test_wheel_package_qnn.sh

@@ -98,7 +98,7 @@ PYTHON_VERSION=$1
 # Check wheel does NOT contain qualcomm/sdk
 # ----------------------------
 echo "Checking wheel does not contain qualcomm/sdk..."
-SDK_FILES=$(unzip -l "$WHEEL_FILE" | awk '{print $4}' | grep "executorch/backends/qualcomm/sdk" || true)
+SDK_FILES=$(unzip -l "$WHEEL_FILE" | awk '{print $4}' | grep -E "executorch/backends/qualcomm/sdk" || true)
 if [ -n "$SDK_FILES" ]; then
     echo "ERROR: Wheel package contains unexpected qualcomm/sdk files:"
     echo "$SDK_FILES"
@@ -111,7 +111,7 @@ fi
 # Check .so files in the wheel
 # ----------------------------
 echo "Checking for .so files inside the wheel..."
-WHEEL_SO_FILES=$(unzip -l "$WHEEL_FILE" | awk '{print $4}' | grep "executorch/backends/qualcomm/python" || true)
+WHEEL_SO_FILES=$(unzip -l "$WHEEL_FILE" | awk '{print $4}' | grep -E "executorch/backends/qualcomm/python" || true)
 if [ -z "$WHEEL_SO_FILES" ]; then
     echo "ERROR: No .so files found in wheel under executorch/backends/qualcomm/python"
     exit 1
@@ -139,8 +139,30 @@ run_core_tests () {
   echo "=== [$LABEL] Installing wheel & deps ==="
   "$PIPBIN" install --upgrade pip
   "$PIPBIN" install "$WHEEL_FILE"
-  "$PIPBIN" install torch=="2.9.0.dev20250906" --index-url "https://download.pytorch.org/whl/nightly/cpu"
-  "$PIPBIN" install --pre torchao --index-url "https://download.pytorch.org/whl/nightly/cpu"
+  TORCH_VERSION=$(
+  "$PYBIN" - <<'PY'
+import runpy
+module_vars = runpy.run_path("torch_pin.py")
+print(module_vars["TORCH_VERSION"])
+PY
+)
+
+  NIGHTLY_VERSION=$(
+  "$PYBIN" - <<'PY'
+import runpy
+module_vars = runpy.run_path("torch_pin.py")
+print(module_vars["NIGHTLY_VERSION"])
+PY
+)
+  echo "=== [$LABEL] Install torch==${TORCH_VERSION}.${NIGHTLY_VERSION} ==="
+
+  # Install torchao based on the pinned PyTorch version
+  "$PIPBIN" install torch=="${TORCH_VERSION}.${NIGHTLY_VERSION}" --index-url "https://download.pytorch.org/whl/nightly/cpu"
+
+  # Install torchao based on the pinned commit from third-party/ao submodule
+  pushd "$REPO_ROOT/third-party/ao" > /dev/null
+  USE_CPP=0 "$PYBIN" setup.py develop
+  popd > /dev/null
 
   echo "=== [$LABEL] Import smoke tests ==="
   "$PYBIN" -c "import executorch; print('executorch imported successfully')"

backends/aoti/utils.h

@@ -36,6 +36,8 @@ inline executorch::aten::ScalarType dtype_to_scalar_type(int32_t dtype) {
   switch (dtype) {
     case 6: // PyTorch's float32 dtype code
       return executorch::aten::ScalarType::Float;
+    case 15: // PyTorch's bfloat16 dtype code
+      return executorch::aten::ScalarType::BFloat16;
     // Future support for additional dtypes can be added here
     default:
       ET_LOG(Error, "Unsupported dtype: %d for ScalarType conversion", dtype);

backends/apple/coreml/compiler/coreml_preprocess.py

@@ -6,6 +6,7 @@
 import logging
 
 import shutil
+import tempfile
 import uuid
 from dataclasses import asdict, dataclass
 from enum import Enum
@@ -415,7 +416,7 @@ def preprocess_model(
         mlmodel: ct.models.MLModel, model_type: MODEL_TYPE
     ) -> PreprocessResult:
         identifier = "executorch_" + str(uuid.uuid4())
-        dir_path: Path = Path("tmp") / identifier
+        dir_path: Path = Path(tempfile.gettempdir()) / identifier
         model_dir_path: Path = dir_path / "lowered_module"
         model_spec: ct.proto.Model_pb2 = mlmodel.get_spec()
         logger.warning(

backends/arm/_passes/arm_pass_manager.py

@@ -112,6 +112,8 @@
 from executorch.exir.pass_manager import PassManager
 from executorch.exir.passes.remove_graph_asserts_pass import RemoveGraphAssertsPass
 from torch.fx import GraphModule
+from torch.fx.passes.infra.pass_base import PassResult
+from torch.nn.modules import Module
 
 
 class ArmPassManager(PassManager):
@@ -355,3 +357,20 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
             self.add_pass(DecomposeMaskedFill())
 
         return self._transform(graph_module)
+
+    def __call__(self, module: Module) -> PassResult:
+        try:
+            return super().__call__(module)
+        except Exception as e:
+            first_exception = e.__cause__ or e.__context__ or e
+            import re
+
+            message = e.args[0]
+            m = re.search(r"An error occurred when running the '([^']+)' pass", message)
+            if m:
+                pass_name = m.group(1)
+                first_exception.args = (
+                    f"{pass_name}: {first_exception.args[0]}",
+                    *first_exception.args[1:],
+                )
+            raise first_exception

backends/arm/_passes/fuse_batchnorm2d_pass.py

@@ -12,6 +12,7 @@
     create_node,
     get_first_fake_tensor,
 )
+from executorch.backends.arm.common.debug import get_node_debug_info
 from executorch.backends.transforms.utils import (
     create_constant_placeholder,
     delete_constant_placeholder,
@@ -60,8 +61,16 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
             input_node = node.all_input_nodes[0]
             is_single_user = len(input_node.users) == 1
             bn_weight_node, bn_bias_node, bn_mean_node, bn_var_node = node.args[1:5]
-            assert bn_mean_node is not None, "Batchnorm mean node cannot be None."
-            assert bn_var_node is not None, "Batchnorm var node cannot be None."
+            if bn_mean_node is None:
+                raise RuntimeError(
+                    "BatchNorm mean buffer missing for node: "
+                    f"{get_node_debug_info(node, graph_module)}"
+                )
+            if bn_var_node is None:
+                raise RuntimeError(
+                    "BatchNorm variance buffer missing for node: "
+                    f"{get_node_debug_info(node, graph_module)}"
+                )
 
             epsilon = node.args[-1]
 
@@ -133,14 +142,23 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
                     input_node = new_input_node
             else:
                 input_weight_node, input_bias_node = input_node.args[1:3]
-                assert (
+                if not (
                     isinstance(input_weight_node, Node)
                     and input_weight_node.op == "placeholder"
-                ), "Parameter weight of convolution must be a placeholder"
-                assert (input_bias_node is None) or (
-                    isinstance(input_weight_node, Node)
-                    and input_weight_node.op == "placeholder"
-                ), "Parameter bias of convolution must be a placeholder or None"
+                ):
+                    raise RuntimeError(
+                        "Parameter weight of convolution must be a placeholder"
+                    )
+                if not (
+                    (input_bias_node is None)
+                    or (
+                        isinstance(input_weight_node, Node)
+                        and input_weight_node.op == "placeholder"
+                    )
+                ):
+                    raise RuntimeError(
+                        "Parameter bias of convolution must be a placeholder or None"
+                    )
 
                 input_weight_tensor = torch.Tensor(
                     get_param(self.exported_program, input_weight_node)

backends/arm/arm_vela.py

@@ -34,7 +34,10 @@ def vela_bin_pack_io(prefix, data):
         io_elem_size = data[prefix + "_elem_size"][i]
         io_offset = data[prefix + "_offset"][i]
         io_region = data[prefix + "_region"][i]
-        assert len(io_shape) == vela_io_shape_dims
+        if len(io_shape) != vela_io_shape_dims:
+            raise ValueError(
+                f"Expected {vela_io_shape_dims}D shape, got {len(io_shape)}D"
+            )
         inp_pad = io_shape.tolist()
         io_struct = struct.pack(
             "<iiiiiiiii", *inp_pad, io_elem_size, io_offset, io_region

backends/arm/common/arm_compile_spec.py

@@ -126,7 +126,8 @@ def validate(self):
 
     def to_list(self):
         """Get the ArmCompileSpec in list form."""
-        assert self.tosa_spec
+        if not self.tosa_spec:
+            raise ValueError("tosa_spec must be set before calling to_list()")
 
         # Always supply a TOSA version
         compile_spec = [

backends/arm/operator_support/ethos_u55_support.py

@@ -384,3 +384,63 @@ def is_node_supported(
             return False
 
         return True
+
+
+class EthosU55CastCheck(OperatorSupportBase):
+    """Reject unsupported casts on U55.
+
+    U55 does not support casting from INT32 or any casts involving BOOL. Note that
+    casting from one dtype to the same dtype is a no-op and is supported.
+
+
+    Attributes:
+        reporter (WhyNoPartitionReporter): Reporter for rejection reasons.
+
+    """
+
+    targets = [
+        exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
+    ]
+
+    def __init__(self, reporter: WhyNoPartitionReporter):
+        """Initialize the check with a reporter.
+
+        Args:
+            reporter (WhyNoPartitionReporter): Reporter for rejection reasons.
+
+        """
+        super().__init__()
+        self.reporter = reporter
+
+    def is_node_supported(
+        self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node
+    ) -> bool:
+        """Return True if the node satisfies the cast constraints of U55.
+
+        Args:
+            submodules (typing.Mapping[str, torch.nn.Module]): Exported modules.
+            node (fx.Node): FX node to check.
+
+        Returns:
+            bool: True if supported; otherwise, False.
+
+        """
+        if node.target not in self.targets:
+            return True
+        input_dtype = get_first_fake_tensor(node.all_input_nodes[0]).dtype
+        output_dtype = get_first_fake_tensor(node).dtype
+        if input_dtype == output_dtype:
+            # This is ok as this will not result in a cast
+            return True
+        if input_dtype in (torch.bool, torch.int32):
+            self.reporter.report_reject(
+                node, f"Casting from {input_dtype} is not supported on U55."
+            )
+            return False
+        if output_dtype in (torch.bool,):
+            self.reporter.report_reject(
+                node, f"Casting to {output_dtype} is not supported on U55."
+            )
+            return False
+
+        return True

backends/arm/operator_support/tosa_supported_operators.py

@@ -21,6 +21,7 @@
 from executorch.backends.arm._passes.insert_table_ops import TableOps
 from executorch.backends.arm.constants import DQ_OPS, Q_OPS
 from executorch.backends.arm.operator_support.ethos_u55_support import (
+    EthosU55CastCheck,
     EthosU55DtypeSupport,
     EthosU55NotSupported,
     EthosU55TransposeCheck,
@@ -141,6 +142,7 @@ def tosa_support_factory(
         negative_checks.append(EthosU55DtypeSupport(reporter))
         negative_checks.append(EthosU55TransposeCheck(reporter))
         negative_checks.append(EthosU55ViewCheck(reporter))
+        negative_checks.append(EthosU55CastCheck(reporter))
 
     return chain(
         reporter.wrap_check(

backends/arm/process_node.py

@@ -106,11 +106,16 @@ def process_inputs_to_parameters(
         ) from e
     parameter_data = get_param(edge_program, node)
 
-    assert isinstance(parameter_data, torch.Tensor), "Expect Attr to be tensor"
+    if not isinstance(parameter_data, torch.Tensor):
+        raise TypeError(
+            f"Expected parameter '{node.name}' to be a torch.Tensor, got "
+            f"{type(parameter_data).__name__}"
+        )
     parameter_values = parameter_data.detach().numpy()
 
     if tosa_arg.dtype == torch.float32:
-        assert tosa_spec.support_float(), f"{tosa_spec} doesn't support float"
+        if not tosa_spec.support_float():
+            raise ValueError(f"{tosa_spec} doesn't support float operations")
 
     # Handle special case for INT48 tensors
     special_type = node.meta.get(TosaSpecialDtype.meta_key(), None)
@@ -142,7 +147,11 @@ def process_inputs_to_buffers(
         ) from e
     buffer_data = get_buffer(edge_program, node)
 
-    assert isinstance(buffer_data, torch.Tensor), "Expect Attr to be tensor"
+    if not isinstance(buffer_data, torch.Tensor):
+        raise TypeError(
+            f"Expected buffer '{node.name}' to be a torch.Tensor, got "
+            f"{type(buffer_data).__name__}"
+        )
     buffer_values = buffer_data.detach().numpy()
 
     # TODO: fragile code for temporary fix
@@ -183,8 +192,12 @@ def process_placeholder(
     tosa_spec: TosaSpecification,
 ):
     """Wrapper for processing and serializing all types of placeholders"""
-    assert node.name == node.target, "Expect placeholder name and target to match"
-    assert 0 == len(node.args), "Can't handle default input values"
+    if node.name != node.target:
+        raise ValueError(
+            f"Placeholder name '{node.name}' does not match target '{node.target}'"
+        )
+    if len(node.args) != 0:
+        raise ValueError(f"Placeholder '{node.name}' must not have default values")
 
     if node.name in edge_program.graph_signature.user_inputs:
         process_inputs(node, tosa_graph, tosa_spec)