shoumikhin · Nov 2, 2026 · Oct 30, 2026 · Oct 31, 2026 · Oct 31, 2026 · Oct 31, 2026
@@ -25,7 +25,6 @@
     is_per_tensor,
     is_qparam,
     is_quant,
-    tag_as_implicit_q_dq,
 )
 from executorch.backends.xnnpack.utils.utils import (
     get_input_node,
@@ -139,11 +138,6 @@ def get_deps(
         valid_deps = valid_bias and valid_weight and valid_act and valid_output
         deps = list(chain(bias_deps, weight_deps, act_deps, output_deps))
 
-        # Tag q/dq nodes as implicit q/dq nodes
-        for dep in deps:
-            if is_dequant(dep) or is_quant(dep):
-                tag_as_implicit_q_dq(dep)
-
         return valid_deps, deps
 
     def _get_weight_deps(
@@ -224,9 +218,12 @@ def _get_output_deps(
             node_users = list(node.users.keys())
             if len(node_users) == 1:
                 n_output = node_users[0]
+                # After a decomposed linear, the activation can have more dims than
+                # XNNPACK allows even when the 2-D matmul does not. Leave it out then.
                 if (
                     n_output.op == "call_function"
                     and format_target_name(n_output.target.__name__) in self.fused_acts
+                    and self._check_tensor_ranks(n_output)
                 ):
                     gemm_deps.append(n_output)
 
@@ -570,17 +567,23 @@ def find_partition_args(input_node):
         node.args = tuple(fake_args)
         node.users = src_partition.output_nodes[0].users
         valid_deps, deps = super().get_deps(node, ep)
+        _, precision = self._overwrite_precision(node)
 
         # Reset addmm node back to old args and users
         node.args = old_args
         node.users = old_users
 
-        # When using force_non_static_weights_for_f32_linear, we want to get_deps to overwrite the source partition nodes.
-        # Else we want to be greedy.
+        # A float linear with force_non_static_weights_for_f32_linear takes only the
+        # deps inside the source partition. Else we want to be greedy, but leave out
+        # views XNNPACK cannot hold.
         ret_deps = (
             list(set(deps) & set(src_partition.nodes))
             if self.force_non_static_weights_for_f32_linear
-            else list(set(deps) | set(src_partition.nodes))
+            and precision == ConfigPrecisionType.FP32
+            else list(
+                set(deps)
+                | {n for n in src_partition.nodes if self._check_tensor_ranks(n)}
+            )
         )
 
         return valid_deps, ret_deps

@@ -16,11 +16,7 @@
     ConfigPrecisionType,
     XNNPartitionerConfig,
 )
-from executorch.backends.xnnpack.utils.quant_utils import (
-    is_dequant,
-    is_quant,
-    tag_as_implicit_q_dq,
-)
+from executorch.backends.xnnpack.utils.quant_utils import is_dequant, is_quant
 from executorch.backends.xnnpack.utils.utils import (
     get_input_node,
     get_param_tensor,
@@ -165,15 +161,6 @@ def get_node_and_deps(
                 # Expected node --> fused_act (optional) --> dequant
                 return deps
 
-            # Tag input nodes (dq nodes) as implicit q/dq nodes
-            for dq_input in node.all_input_nodes:
-                if is_dequant(dq_input):
-                    tag_as_implicit_q_dq(dq_input)
-
-            # Tag node_output (q node) as an implicit q/dq node
-            if is_quant(node_output):
-                tag_as_implicit_q_dq(node_output)
-
             quantized_deps.append(node_output)
 
         return deps + quantized_deps

@@ -14,7 +14,9 @@
     is_dequant,
     is_qparam,
     is_quant,
+    tag_as_implicit_q_dq,
 )
+from executorch.backends.xnnpack.utils.xnnpack_constants import XNN_MAX_TENSOR_DIMS
 from executorch.exir.backend.canonical_partitioners.config_partitioner import (
     format_target_name,
     PartitionerConfig,
@@ -66,14 +68,31 @@ def get_partition(
         with the name of the XNNPartitionerConfig class which return this set of nodes.
         This enforces that all partitions returned by XNNPartitioner configs are labeled
         with the partitioner config which returned them
+
+        Returns an empty list if any of the nodes has more dims than XNNPACK supports.
+        Otherwise marks every quantize and dequantize node in the partition, other
+        than the given node, as implicit, so that it is folded into the op it feeds
+        or follows.
         """
         partitioned_nodes = self.get_node_and_deps(node, ep)
+        # A partition can take nodes that their own configs reject.
+        for n in partitioned_nodes:
+            if not self._check_tensor_ranks(n):
+                why(node, reason=f"its partition takes {n}, which has too many dims")
+                return []
+
+        # Fold only now, so a rejected node keeps its q/dq nodes as separate ops. A
+        # q/dq node partitioned on its own stays a separate op as well.
+        for n in partitioned_nodes:
+            if n is not node and (is_quant(n) or is_dequant(n)):
+                tag_as_implicit_q_dq(n)
+
         # label partitioned nodes with the name of the partitioner config
-        for node in partitioned_nodes:
-            if "xnn_partitioner_config" in node.meta:
-                node.meta["xnn_partitioner_config"].append(self.__class__.__name__)
+        for n in partitioned_nodes:
+            if "xnn_partitioner_config" in n.meta:
+                n.meta["xnn_partitioner_config"].append(self.__class__.__name__)
             else:
-                node.meta["xnn_partitioner_config"] = [self.__class__.__name__]
+                n.meta["xnn_partitioner_config"] = [self.__class__.__name__]
 
         return partitioned_nodes
 
@@ -94,7 +113,9 @@ def get_node_and_deps(
     ) -> List[torch.fx.Node]:
         """
         Takes in a node and its exported program and returns a list of nodes
-        and its dependencies that need to be partitioned together
+        and its dependencies that need to be partitioned together. Quantize and
+        dequantize nodes in the list, other than the node itself, are folded into
+        the partition's op once get_partition takes it.
 
         Args:
             node: Node to be partitioned
@@ -152,6 +173,20 @@ def check_common_constraints(
             why(node, reason="invalid dtype")
             return False
 
+        return self._check_tensor_ranks(node)
+
+    def _check_tensor_ranks(self, node: torch.fx.Node) -> bool:
+        vals = [arg.meta.get("val") for arg in node.all_input_nodes]
+        node_val = node.meta.get("val")
+        vals.extend(node_val if isinstance(node_val, (list, tuple)) else [node_val])
+        for val in vals:
+            if isinstance(val, torch.Tensor) and val.dim() > XNN_MAX_TENSOR_DIMS:
+                why(
+                    node,
+                    reason=f"XNNPACK supports at most {XNN_MAX_TENSOR_DIMS} dims, got {val.dim()}",
+                )
+                return False
+
         return True
 
     def _check_inputs_are_valid_dtypes(self, node, valid_dtypes):

@@ -257,8 +257,6 @@ class QuantPattern:
     torch.ops.aten.conv_transpose1d.default,
     torch.ops.aten.conv_transpose2d,
     torch.ops.aten.conv_transpose2d.input,
-    torch.ops.aten.conv_transpose3d,
-    torch.ops.aten.conv_transpose3d.input,
 }
 
 LINEAR_TARGETS = {

@@ -240,6 +240,11 @@ def _annotate_linear_relu(
     return annotated_partitions
 
 
+def _is_1d_or_2d_conv(weight_shape: Optional[torch.Size]) -> bool:
+    # XNNPACK runs no other convs.
+    return weight_shape is None or len(weight_shape) <= 4
+
+
 def _do_annotate_conv(
     gm: torch.fx.GraphModule,
     quantization_config: Optional[QuantizationConfig],
@@ -274,9 +279,11 @@ def _do_annotate_conv(
         assert isinstance(weight, Node)
         weight_qspec = get_weight_qspec(quantization_config)
         num_groups = get_groups_from_conv(conv_node)
+        weight_shape = getattr(weight.meta.get("val", None), "shape", None)
 
         # skip if transposed conv has more than 1 group
         skip = skip or (is_conv_transpose and num_groups != 1)
+        skip = skip or not _is_1d_or_2d_conv(weight_shape)
 
         if is_conv_transpose:
             # transposed convs per output channel quantization
@@ -291,8 +298,6 @@ def _do_annotate_conv(
 
         # Only annotate dynamically quantized conv if it's 2D and not depthwise
         if is_dynamic:
-            weight_val = weight.meta.get("val", None)
-            weight_shape = getattr(weight_val, "shape", None)
             # Skip if not a 4D weight tensor (i.e. not conv2d)
             skip = skip or (weight_shape is not None and len(weight_shape) != 4)
             # Skip if depthwise (default to groups=1 since it's not an arg)
@@ -356,6 +361,7 @@ def _do_annotate_conv_relu(
         assert isinstance(weight, Node)
         weight_qspec = get_weight_qspec(quantization_config)
         groups = get_groups_from_conv(conv_node)
+        weight_shape = getattr(weight.meta.get("val", None), "shape", None)
         if is_conv_transpose:
             # transposed convs per output channel quantization
             weight_qspec = change_quantization_config(weight_qspec, ch_axis=1)
@@ -371,7 +377,7 @@ def _do_annotate_conv_relu(
         if _is_annotated(partition):
             continue
 
-        if is_conv_transpose and groups != 1:
+        if (is_conv_transpose and groups != 1) or not _is_1d_or_2d_conv(weight_shape):
             continue
 
         if filter_fn and not filter_fn(conv_node):

@@ -96,6 +96,7 @@ fbcode_target(_kind = runtime.python_test,
         "test_xnnpack_utils.py",
     ],
     supports_static_listing = False,
+    preload_deps = ["//executorch/kernels/quantized:aot_lib"],
     deps = [
         "fbsource//third-party/pypi/parameterized:parameterized",
         "//executorch/backends/test/harness:tester",

@@ -6,6 +6,7 @@
 
 import itertools
 import unittest
+from dataclasses import replace
 from typing import Optional
 
 import torch
@@ -355,6 +356,34 @@ def test_fp32_conv3d_single_element_stride_doesnt_partition(self) -> None:
             .run_method_and_compare_outputs()
         )
 
+    def test_qs8_quantizer_skips_conv3d(self) -> None:
+        # XNNPACK has no conv3d support, so its quantizer leaves 3d convs in float,
+        # with or without a fused relu. The same model in 2d gets its activations
+        # and weights quantized. The weights are quantized ahead of time, so only
+        # their dequantize nodes remain.
+        quant_config = get_symmetric_quantization_config(is_per_channel=True)
+        ops = "torch.ops.quantized_decomposed"
+        for dims, conv, conv_transpose, quantized_activations, quantized_weights in (
+            (2, torch.nn.Conv2d, torch.nn.ConvTranspose2d, 4, 3),
+            (3, torch.nn.Conv3d, torch.nn.ConvTranspose3d, 0, 0),
+        ):
+            m = torch.nn.Sequential(
+                conv(2, 2, 2), torch.nn.ReLU(), conv(2, 2, 2), conv_transpose(2, 2, 2)
+            )
+            (
+                Tester(m.eval(), (torch.randn([1, 2] + [4] * dims),))
+                .quantize(Quantize(quantization_config=quant_config))
+                .export()
+                .check_count(
+                    {
+                        f"{ops}.quantize_per_tensor.default": quantized_activations,
+                        f"{ops}.dequantize_per_tensor.default": quantized_activations,
+                        f"{ops}.quantize_per_channel.default": 0,
+                        f"{ops}.dequantize_per_channel.default": quantized_weights,
+                    }
+                )
+            )
+
     def test_fp32_conv2d_permute(self) -> None:
         for transpose in (True, False):
             for perm_order in list(itertools.permutations([0, 1, 2, 3])):
@@ -375,6 +404,26 @@ def test_qs8_conv2d_per_channel(self) -> None:
                 quant_config=get_symmetric_quantization_config(is_per_channel=True),
             )
 
+    def test_qs8_conv2d_float_output(self) -> None:
+        # XNNPACK leaves out a quantized conv whose output is not quantized. The
+        # quantize and dequantize of its input then run in a delegate of their own.
+        quant_config = replace(
+            get_symmetric_quantization_config(is_per_channel=True),
+            output_activation=None,
+        )
+        m = Conv2d()
+        (
+            Tester(m.eval(), m.get_inputs())
+            .quantize(Quantize(quantization_config=quant_config))
+            .export()
+            .to_edge_transform_and_lower()
+            .check(["executorch_exir_dialects_edge__ops_aten_convolution_default"])
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs()
+        )
+
     def test_qs8_conv2d_even_kernel_same_padding(self) -> None:
         # An even-kernel 'same'-padding conv decomposes into
         # dequant -> constant_pad_nd -> convolution. The pad and conv are pulled
@@ -832,14 +881,27 @@ def get_inputs(self):
                         conv_count=1,
                     )
                 else:
-                    with self.assertRaises(RuntimeError):
-                        self._test(
-                            PerChannelConvTranspose2d(
-                                3 * groups, 5 * groups, groups, ch_axis
-                            ),  # ch_axis=0
-                            quant_config=None,
-                            conv_count=1,
+                    # XNNPACK leaves this conv out, and the quantize and dequantize
+                    # nodes around it then run in delegates of their own.
+                    m = PerChannelConvTranspose2d(
+                        3 * groups, 5 * groups, groups, ch_axis
+                    )
+                    (
+                        Tester(m.eval(), m.get_inputs())
+                        .export()
+                        .to_edge_transform_and_lower()
+                        .check(
+                            [
+                                "executorch_exir_dialects_edge__ops_aten_convolution_default"
+                            ]
                         )
+                        .check_count(
+                            {"torch.ops.higher_order.executorch_call_delegate": 2}
+                        )
+                        .to_executorch()
+                        .serialize()
+                        .run_method_and_compare_outputs(qtol=1)
+                    )
 
     def test_fp32_tconv_output_padding(self):
         """Test transposed convolution with non-zero output padding."""