XNNPACK: fold constant subgraphs before ATen decomposition

backends/xnnpack/_passes/BUCK

@@ -18,6 +18,7 @@ fbcode_target(_kind = runtime.python_library,
         "//executorch/exir:pass_base",
         "//executorch/exir/dialects:lib",
         "//executorch/exir/passes:const_prop_pass",
+        "//executorch/exir/passes:constant_prop_pass",
         "//executorch/exir/passes:memory_format_ops_pass",
         "//executorch/exir/program:program",
         "//executorch/backends/transforms:utils",

backends/xnnpack/_passes/fold_constant_subgraphs.py

@@ -0,0 +1,155 @@
+# Copyright 2026 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 operator
+
+import torch
+
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.passes.constant_prop_pass import constant_prop_pass
+from torch.export import ExportedProgram
+
+_CONSTANT_SUBGRAPH_OPS = {
+    torch.ops.aten.cat.default,
+    torch.ops.aten.chunk.default,
+    torch.ops.aten.split.Tensor,
+    torch.ops.aten.split_with_sizes.default,
+    operator.getitem,
+    torch.ops.aten._weight_norm.default,
+}
+
+
+def _has_unsafe_tuple_consumer(exported_program: ExportedProgram) -> bool:
+    """Avoid propagating tuple-valued chunk/split nodes, which
+    constant_prop_pass cannot replace."""
+    tuple_producer_targets = {
+        torch.ops.aten.chunk.default,
+        torch.ops.aten.split.Tensor,
+        torch.ops.aten.split_with_sizes.default,
+    }
+    for node in exported_program.graph.nodes:
+        if (
+            node.op == "call_function"
+            and node.target in tuple_producer_targets
+            and any(
+                user.op != "call_function" or user.target is not operator.getitem
+                for user in node.users
+            )
+        ):
+            return True
+    return False
+
+
+def _has_unsafe_mutation(node: torch.fx.Node, schema) -> bool:
+    for index, argument in enumerate(schema.arguments):
+        if argument.alias_info is None or not argument.alias_info.is_write:
+            continue
+        mutated = (
+            node.args[index]
+            if index < len(node.args)
+            else node.kwargs.get(argument.name)
+        )
+        producer = mutated if isinstance(mutated, torch.fx.Node) else None
+        producer_schema = (
+            getattr(producer.target, "_schema", None)
+            if producer is not None and producer.op == "call_function"
+            else None
+        )
+        if (
+            producer_schema is None
+            or producer_schema.is_mutable
+            or producer.target in _CONSTANT_SUBGRAPH_OPS
+            or any(ret.alias_info is not None for ret in producer_schema.returns)
+        ):
+            return True
+    return False
+
+
+def _has_unsafe_node_effect(node: torch.fx.Node) -> bool:
+    if isinstance(node.target, torch._ops.HigherOrderOperator):
+        return True
+    schema = getattr(node.target, "_schema", None)
+    if schema is not None and schema.is_mutable:
+        return _has_unsafe_mutation(node, schema)
+    return node.is_impure() and node.target in _CONSTANT_SUBGRAPH_OPS
+
+
+def _has_unsafe_graph_effects(exported_program: ExportedProgram) -> bool:
+    """Return whether graph effects make constant folding unsafe."""
+    return any(
+        node.op == "call_function" and _has_unsafe_node_effect(node)
+        for node in exported_program.graph.nodes
+    )
+
+
+class FoldConstantSubgraphsPass:
+    """
+    Fold selected constant subgraphs into direct constants before decomposition.
+
+    This pass folds constant-only subgraphs such as:
+
+    bias_0  bias_1  bias_2       # constant parameters
+     \\    ||      //
+         cat
+          |
+      split/chunk
+          |
+        getitem
+          |
+     addmm / linear
+
+    Although the resulting bias slices are constant, the intermediate
+    cat/split/getitem chain can prevent later lowering and delegation from
+    recognizing them as constant inputs.
+
+    Constant propagation evaluates eligible operations ahead of time and replaces
+    each projection with a direct constant:
+
+    constant bias_0 -> addmm / linear
+    constant bias_1 -> addmm / linear
+    constant bias_2 -> addmm / linear
+
+    Folding is skipped when mutation, aliasing, impurity, or higher-order
+    operations could make replacing a runtime value with persistent constant
+    data unsafe.
+    """
+
+    def __call__(self, exported_program: ExportedProgram) -> ExportedProgram:
+        # This runs before functionalization, so inspect schemas as well as the
+        # graph signature for mutation.
+        if any(
+            "MUTATION" in str(spec.kind)
+            for spec in exported_program.graph_signature.output_specs
+        ):
+            return exported_program
+        if _has_unsafe_graph_effects(exported_program):
+            return exported_program
+
+        if not any(
+            node.op == "call_function" and node.target in _CONSTANT_SUBGRAPH_OPS
+            for node in exported_program.graph.nodes
+        ):
+            return exported_program
+
+        if _has_unsafe_tuple_consumer(exported_program):
+            return exported_program
+        skipped = {
+            node.target
+            for node in exported_program.graph.nodes
+            if node.op == "call_function" and node.target not in _CONSTANT_SUBGRAPH_OPS
+        }
+        folded_program = constant_prop_pass(
+            exported_program, custom_skip_targets=skipped
+        )
+        # Constant propagation inserts an Edge clone for constant outputs, but
+        # pre-decomposition graphs must remain entirely in the ATen dialect.
+        replaced_clone = False
+        for node in folded_program.graph.nodes:
+            if node.target == exir_ops.edge.aten.clone.default:
+                node.target = torch.ops.aten.clone.default
+                replaced_clone = True
+        if replaced_clone:
+            folded_program.graph_module.recompile()
+        return folded_program

backends/xnnpack/operators/node_visitor.py

@@ -1,24 +1,21 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2026 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 ctypes
 import hashlib
 import logging
-
 from typing import cast, Dict, List, Optional, Tuple
 
 import torch
 from executorch.backends.transforms import get_shape
-
 from executorch.backends.xnnpack._passes.channels_last_tagged_reshape_pass import (
     ChannelsLastTaggedReshapePass,
 )
-
 from executorch.backends.xnnpack.operators.quant_params import QuantParams
-
 from executorch.backends.xnnpack.serialization.xnnpack_graph_schema import (
     ConstantDataOffset,
     PerChannelGroupQuant,
@@ -39,7 +36,6 @@
     is_param_node,
     PERM_NCHW_TO_NHWC,
 )
-
 from executorch.backends.xnnpack.utils.xnnpack_constants import (
     UINT64_MAX,
     XNN_INVALID_VALUE_ID,
@@ -80,6 +76,17 @@ def get_tensor_value(xvalue: XValue) -> XNNTensorValue:
         return q_tensor.tensor_value
 
 
+def _tensor_bytes(value: torch.Tensor) -> bytes:
+    """Return the tensor's contiguous physical storage as bytes."""
+
+    assert value.is_contiguous() or value.is_contiguous(
+        memory_format=torch.channels_last
+    ), "Tensor must have contiguous storage"
+
+    # data_ptr() accounts for storage_offset(), limiting the read to this tensor.
+    return ctypes.string_at(value.data_ptr(), value.numel() * value.element_size())
+
+
 class NodeVisitor:
     """
     Node visitor pattern for visiting nodes in an edge IR graph and
@@ -288,16 +295,9 @@ def get_quant_params(
             if quant_params.per_channel_group:
                 scale = scale.to(torch.bfloat16)
 
-            # The bytes are read from the start of the storage in memory order.
-            # Scales are small, so always copy them into storage of their own.
-            scale = scale.clone(memory_format=torch.contiguous_format)
-
-            num_bytes = scale.untyped_storage().nbytes()
-            scale_array = ctypes.cast(
-                scale.untyped_storage().data_ptr(),
-                ctypes.POINTER(ctypes.c_char * num_bytes),
-            ).contents
-            scale_name = hashlib.sha256(bytes(scale_array)).hexdigest()
+            scale_data = _tensor_bytes(scale.contiguous())
+            num_bytes = len(scale_data)
+            scale_name = hashlib.sha256(scale_data).hexdigest()
             scale_name = "scale_" + scale_name
             xnn_graph.constant_data.append(
                 ConstantDataOffset(
@@ -309,7 +309,7 @@ def get_quant_params(
                     f"Adding constant data with name, key {scale_name} and external_tag {external_tag} to named_data_store"
                 )
             self._named_data_store.add_named_data(
-                scale_name, bytes(scale_array), CONSTANT_TENSOR_ALIGNMENT, external_tag
+                scale_name, scale_data, CONSTANT_TENSOR_ALIGNMENT, external_tag
             )
 
             if quant_params.per_channel_group:
@@ -623,26 +623,16 @@ def get_serialized_buffer_index(
         if quant_params is not None and quant_params.is_qc4w:
             const_val = self.convert_to_qc4w(const_val)
 
-        # The bytes are read from the start of the storage, so a view of a
-        # larger tensor needs its own copy.
-        if const_val.nbytes != const_val.untyped_storage().nbytes():
-            const_val = const_val.clone()
-
-        size = const_val.untyped_storage().nbytes()
-        array_type = ctypes.c_char * size
-        array = ctypes.cast(
-            const_val.untyped_storage().data_ptr(),
-            ctypes.POINTER(array_type),
-        ).contents
+        data = _tensor_bytes(const_val)
+        size = len(data)
 
         check_or_raise(
             size > 0,
             f"Serializing constant data node {tensor} but tensor value has no bytes",
         )
-        sha256_hash = hashlib.sha256(bytes(array))
+        sha256_hash = hashlib.sha256(data)
         named_key = sha256_hash.hexdigest()
 
-        size = const_val.untyped_storage().nbytes()
         xnn_graph.constant_data.append(
             ConstantDataOffset(offset=UINT64_MAX, size=size, named_key=named_key)
         )
@@ -657,7 +647,7 @@ def get_serialized_buffer_index(
             )
         self._named_data_store.add_named_data(
             named_key,
-            bytes(array),
+            data,
             alignment=CONSTANT_TENSOR_ALIGNMENT,
             external_tag=external_tag,
         )

backends/xnnpack/partition/BUCK

@@ -12,6 +12,7 @@ fbcode_target(_kind = runtime.python_library,
     deps = [
         ":partitioner_graphs",
         "//executorch/backends/xnnpack:xnnpack_preprocess",
+        "//executorch/backends/xnnpack/_passes:xnnpack_passes",
         "//executorch/backends/xnnpack/partition/config:xnnpack_partitioner_configs",
         "//executorch/exir:delegate",
         "//executorch/exir:lib",

backends/xnnpack/partition/xnnpack_partitioner.py

@@ -1,5 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2026 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.
@@ -9,12 +10,14 @@
 import logging
 from typing import List, Optional, Type, Union
 
+from executorch.backends.xnnpack._passes.fold_constant_subgraphs import (
+    FoldConstantSubgraphsPass,
+)
 from executorch.backends.xnnpack.partition.config import ALL_PARTITIONER_CONFIGS
 from executorch.backends.xnnpack.partition.config.xnnpack_config import (
     ConfigPrecisionType,
     XNNPartitionerConfig,
 )
-
 from executorch.backends.xnnpack.xnnpack_preprocess import XnnpackBackend
 from executorch.exir.backend.backend_details import ExportedProgram
 from executorch.exir.backend.canonical_partitioners.config_partitioner import (
@@ -28,6 +31,12 @@
 
 
 class XnnpackPartitioner(ConfigerationBasedPartitioner):
+    def transform_for_pre_decomposition(
+        self, exported_program: ExportedProgram
+    ) -> ExportedProgram:
+        """Apply XNNPACK transforms before default ATen decompositions."""
+        return FoldConstantSubgraphsPass()(exported_program)
+
     def __init__(
         self,
         configs: Optional[List[Type[XNNPartitionerConfig]]] = None,
@@ -91,6 +100,7 @@ def partition(self, exported_program):
         if self._check_if_called_from_to_backend():
             logger.warning(
                 "\nDEPRECATION WARNING: You are using the deprecated 'to_edge() + to_backend()' workflow. "
+                "It does not run pre-decomposition transforms, including constant cat/split folding. "
                 "Please consider migrating to 'to_edge_transform_and_lower()' for better error handling and optimization. "
             )