[XNNPACK] Lower upsample_nearest2d x2 as concatenate4 + depth_to_space_2d

backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py

@@ -61,6 +61,7 @@ class ChannelsLastTaggedReshapePass(XNNPACKPass):
     memory_sensitive_ops_nhwc = {
         exir_ops.edge.aten.convolution.default,
         exir_ops.edge.aten.upsample_bilinear2d.vec,
+        exir_ops.edge.aten.upsample_nearest2d.vec,
         exir_ops.edge.aten.mean.dim,
         exir_ops.edge.aten.max_pool2d.default,
         exir_ops.edge.aten.amax.default,

backends/xnnpack/operators/__init__.py

@@ -57,5 +57,6 @@
     op_sub,
     op_tanh,
     op_to_copy,
+    op_upsample_nearest2d,
     op_view_copy,
 )

backends/xnnpack/operators/op_upsample_nearest2d.py

@@ -0,0 +1,164 @@
+# Copyright (c) 2026 Software Mansion
+# 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 Dict, List, Optional
+
+import torch
+from executorch.backends.transforms import get_shape
+from executorch.backends.xnnpack.operators.node_visitor import (
+    NodeVisitor,
+    register_node_visitor,
+)
+from executorch.backends.xnnpack.operators.quant_params import QuantParams
+from executorch.backends.xnnpack.serialization.xnnpack_graph_schema import (
+    XNNConcatenate4,
+    XNNConvert,
+    XNNDepthToSpace2d,
+    XNNGraph,
+    XNNQuantizedTensorValue,
+    XNNTensorValue,
+    XNode,
+    XValue,
+)
+from executorch.backends.xnnpack.utils.utils import check_or_raise, get_input_node
+from executorch.backends.xnnpack.utils.xnnpack_constants import XNN_INVALID_VALUE_ID
+
+
+@register_node_visitor
+class UpsampleNearest2dVisitor(NodeVisitor):
+    """
+    XNNPACK has no nearest resize, so x2 upsampling is lowered as four copies
+    of the input concatenated along channels, [N, H, W, 4C], followed by
+    depth_to_space_2d(block_size=2). Neither op changes values, so the
+    intermediates carry the input's quantization parameters, and an output
+    quantized differently is requantized with a convert.
+    """
+
+    target = "aten.upsample_nearest2d.vec"
+
+    def __init__(self, *args) -> None:
+        super().__init__(*args)
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+        xnn_graph: XNNGraph,
+        vals_to_ids: Dict[torch.fx.Node, int],
+        debug_handle: int,
+    ) -> None:
+        input_node = get_input_node(node, 0)
+        input_quant_params = QuantParams.from_inputs(input_node, self._exported_program)
+        output_quant_params = QuantParams.from_outputs(node)
+        self.define_tensor(
+            input_node,
+            xnn_graph,
+            vals_to_ids,
+            convert_to_nhwc=True,
+            quant_params=input_quant_params,
+        )
+        self.define_tensor(
+            node,
+            xnn_graph,
+            vals_to_ids,
+            convert_to_nhwc=True,
+            quant_params=output_quant_params,
+        )
+        input_id = vals_to_ids[input_node]
+        output_id = vals_to_ids[node]
+
+        in_shape = get_shape(input_node)
+        out_shape = get_shape(node)
+        check_or_raise(len(in_shape) == 4, "upsample_nearest2d expects a 4D input")
+        n, c, h, w = in_shape
+        check_or_raise(
+            out_shape[-2] == 2 * h and out_shape[-1] == 2 * w,
+            f"only x2 nearest upsampling is supported, got {in_shape} -> {out_shape}",
+        )
+
+        concat_id = self.define_intermediate(
+            xnn_graph, node, [n, h, w, 4 * c], input_quant_params
+        )
+        if self.quant_params_match(input_quant_params, output_quant_params):
+            d2s_output_id = output_id
+        else:
+            d2s_output_id = self.define_intermediate(
+                xnn_graph, node, [n, 2 * h, 2 * w, c], input_quant_params
+            )
+
+        xnn_graph.xnodes.append(
+            XNode(
+                xnode_union=XNNConcatenate4(
+                    axis=3,
+                    input1_id=input_id,
+                    input2_id=input_id,
+                    input3_id=input_id,
+                    input4_id=input_id,
+                    output_id=concat_id,
+                    flags=0,
+                    input5_id=XNN_INVALID_VALUE_ID,
+                ),
+                debug_handle=debug_handle,
+            )
+        )
+        xnn_graph.xnodes.append(
+            XNode(
+                xnode_union=XNNDepthToSpace2d(
+                    block_size=2,
+                    input_id=concat_id,
+                    output_id=d2s_output_id,
+                    flags=0,
+                ),
+                debug_handle=debug_handle,
+            )
+        )
+        if d2s_output_id != output_id:
+            xnn_graph.xnodes.append(
+                XNode(
+                    xnode_union=XNNConvert(
+                        input_id=d2s_output_id, output_id=output_id, flags=0
+                    ),
+                    debug_handle=debug_handle,
+                )
+            )
+
+    def define_intermediate(
+        self,
+        xnn_graph: XNNGraph,
+        node: torch.fx.Node,
+        dims: List[int],
+        quant_params: Optional[QuantParams],
+    ) -> int:
+        tensor_id = len(xnn_graph.xvalues)
+        tensor_value = XNNTensorValue(
+            datatype=self.get_serialized_dtype(quant_params, node),
+            num_dims=len(dims),
+            dims=dims,
+            external_id=XNN_INVALID_VALUE_ID,
+            constant_buffer_idx=0,
+            flags=0,
+            id_out=tensor_id,
+        )
+        xnn_graph.xvalues.append(
+            XValue(xvalue_union=tensor_value)
+            if quant_params is None
+            else XValue(
+                xvalue_union=XNNQuantizedTensorValue(
+                    tensor_value=tensor_value,
+                    quant_params=self.get_quant_params(quant_params, xnn_graph),
+                )
+            )
+        )
+        return tensor_id
+
+    @staticmethod
+    def quant_params_match(a: Optional[QuantParams], b: Optional[QuantParams]) -> bool:
+        if a is None or b is None:
+            return a is b
+        check_or_raise(
+            not a.per_channel and not b.per_channel,
+            "upsample_nearest2d expects per-tensor quantized activations",
+        )
+        return a.dtype == b.dtype and a.scale == b.scale and a.zp == b.zp

backends/xnnpack/partition/config/__init__.py

@@ -60,6 +60,7 @@
     ToDimOrderCopyConfig,
     UnsqueezeCopyConfig,
     UpsampleBilinear2dConfig,
+    UpsampleNearest2dConfig,
     ViewCopyConfig,
 )
 from executorch.backends.xnnpack.partition.config.node_configs import (
@@ -125,6 +126,7 @@
     SubConfig,
     UnsqueezeCopyConfig,
     UpsampleBilinear2dConfig,
+    UpsampleNearest2dConfig,
     ViewCopyConfig,
     # Quant/Dequant Op Configs
     QuantizedPerTensorConfig,

backends/xnnpack/partition/config/generic_node_configs.py

@@ -27,6 +27,7 @@
     is_param_node,
     normalize_mean_dims,
     normalize_pool2d_args,
+    why_not_x2_nearest_upsample,
 )
 from executorch.exir.backend.canonical_partitioners.config_partitioner import (
     format_target_name,
@@ -503,6 +504,26 @@ def get_original_aten(self) -> Optional[torch._ops.OpOverload]:
         return torch.ops.aten.upsample_bilinear2d.vec
 
 
+class UpsampleNearest2dConfig(GenericNodePartitionerConfig):
+    target_name = "upsample_nearest2d.vec"
+
+    def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
+        if not self.check_common_constraints(node, ep):
+            return False
+
+        reason = why_not_x2_nearest_upsample(node)
+        if reason is not None:
+            why(node, reason=reason)
+            return False
+        return True
+
+    def supported_precision_types(self) -> List[ConfigPrecisionType]:
+        return [ConfigPrecisionType.FP32, ConfigPrecisionType.STATIC_QUANT]
+
+    def get_original_aten(self) -> Optional[torch._ops.OpOverload]:
+        return torch.ops.aten.upsample_nearest2d.vec
+
+
 class ExpConfig(GenericNodePartitionerConfig):
     target_name = "exp.default"
 

backends/xnnpack/quantizer/xnnpack_quantizer.py

@@ -272,6 +272,8 @@ class QuantPattern:
 
 CAT_TARGETS = {torch.ops.aten.cat.default}
 
+UPSAMPLE_NEAREST2D_TARGETS = {torch.ops.aten.upsample_nearest2d.vec}
+
 
 class XNNPACKQuantizer(Quantizer):
     supported_config_and_operators = _get_supported_config_and_operators()
@@ -292,6 +294,7 @@ class XNNPACKQuantizer(Quantizer):
         QuantPattern("mul_relu", False, False, MUL_TARGETS),
         QuantPattern("mul", False, False, MUL_TARGETS),
         QuantPattern("cat", False, False, CAT_TARGETS),
+        QuantPattern("upsample_nearest2d", False, False, UPSAMPLE_NEAREST2D_TARGETS),
     ]
 
     def __init__(self) -> None:

backends/xnnpack/quantizer/xnnpack_quantizer_utils.py

@@ -12,6 +12,7 @@
 from executorch.backends.xnnpack.utils.utils import (
     get_groups_from_conv,
     is_depthwise_conv,
+    why_not_x2_nearest_upsample,
 )
 from torch._subclasses import FakeTensor
 from torch.fx import Node
@@ -727,6 +728,44 @@ def _annotate_adaptive_avg_pool2d(
     return annotated_partitions
 
 
+@register_annotator("upsample_nearest2d")
+def _annotate_upsample_nearest2d(
+    gm: torch.fx.GraphModule,
+    quantization_config: Optional[QuantizationConfig],
+    filter_fn: Optional[Callable[[Node], bool]] = None,
+) -> Optional[list[list[Node]]]:
+    """Nearest upsampling only copies values, so its output shares the
+    input's quantization parameters. Only the x2 case XNNPACK lowers is
+    annotated; any other stays a float op outside the delegate."""
+    annotated_partitions = []
+    for node in gm.graph.nodes:
+        if node.target != torch.ops.aten.upsample_nearest2d.vec:
+            continue
+        if _is_annotated([node]) or (filter_fn and not filter_fn(node)):
+            continue
+        if why_not_x2_nearest_upsample(node) is not None:
+            continue
+
+        input_act = node.args[0]
+        assert isinstance(input_act, Node)
+        if (
+            Q_ANNOTATION_KEY not in input_act.meta
+            or not input_act.meta[Q_ANNOTATION_KEY]._annotated
+            or input_act.meta[Q_ANNOTATION_KEY].output_qspec is None
+        ):
+            input_act_qspec = get_input_act_qspec(quantization_config)
+        else:
+            input_act_qspec = SharedQuantizationSpec(input_act)
+
+        node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation(
+            input_qspec_map={input_act: input_act_qspec},
+            output_qspec=SharedQuantizationSpec((input_act, node)),
+            _annotated=True,
+        )
+        annotated_partitions.append([node])
+    return annotated_partitions
+
+
 def _is_input_large_scalar(node: Node, gm: torch.fx.GraphModule):
     """Check if input is a large scalar value. So that we can skip quantization for the node
     since histc op (in HistogramObserver) only works for values up to certain upper bound

backends/xnnpack/runtime/XNNCompiler.cpp

@@ -1642,6 +1642,39 @@ Error defineConcatenate5Node(
   return Error::Ok;
 }
 
+/*
+Defines serialized depth_to_space_2d node into the subgraph,
+using the remapped ids to map the serialized ids,
+to the new ids generated when defining the tensor value
+*/
+Error defineDepthToSpace2dNode(
+    xnn_subgraph_t subgraph_ptr,
+    const std::unordered_map<uint32_t, uint32_t>& remapped_ids,
+    const NodePtr node,
+    const fb_xnnpack::XNNGraph* graph) noexcept {
+  MAYBE_UNUSED(graph);
+
+  auto graph_node = node->xnode_union_as_XNNDepthToSpace2d();
+
+  REMAP_ID(remapped_ids, graph_node->input_id(), d2s_input);
+  REMAP_ID(remapped_ids, graph_node->output_id(), d2s_output);
+
+  xnn_status status = xnn_define_depth_to_space_2d(
+      subgraph_ptr,
+      graph_node->block_size(),
+      d2s_input,
+      d2s_output,
+      graph_node->flags());
+  ET_CHECK_OR_RETURN_ERROR(
+      status == xnn_status_success,
+      Internal,
+      "Failed to create depth_to_space node %i with code: %s",
+      node->debug_handle(),
+      xnn_status_to_string(status));
+
+  return Error::Ok;
+}
+
 /*
 Defines serialized static_slice node into the subgraph,
 using the remapped ids to map the serialized ids,
@@ -2052,6 +2085,7 @@ DefineNodeFunc getDefineNodeFunc(fb_xnnpack::XNodeUnion nodeType) {
     _DEFINE(Convert)
     _DEFINE(GlobalAvgPooling2d)
     _DEFINE(StaticReshape)
+    _DEFINE(DepthToSpace2d)
     _DEFINE(ArgMaxPooling2d)
     _DEFINE(Concatenate2)
     _DEFINE(Concatenate3)

backends/xnnpack/serialization/runtime_schema.fbs

@@ -160,6 +160,7 @@ union XNodeUnion {
   XNNCopy: _XNNNode1x1,
   XNNCos: _XNNNode1x1,
   XNNApproxGelu: _XNNNode1x1,
+  XNNDepthToSpace2d,
 }
 
 union XValueUnion {
@@ -298,6 +299,13 @@ table XNNStaticReshape {
   flags: uint;
 }
 
+table XNNDepthToSpace2d {
+  block_size: uint;
+  input_id: uint;
+  output_id: uint;
+  flags: uint;
+}
+
 table XNNStaticSlice {
   num_dims:uint;
   offsets:[uint];

backends/xnnpack/serialization/schema.fbs

@@ -156,6 +156,7 @@ union XNodeUnion {
   XNNCopy: _XNNNode1x1,
   XNNCos: _XNNNode1x1,
   XNNApproxGelu: _XNNNode1x1,
+  XNNDepthToSpace2d,
 }
 
 union XValueUnion {
@@ -294,6 +295,13 @@ table XNNStaticReshape {
   flags: uint;
 }
 
+table XNNDepthToSpace2d {
+  block_size: uint;
+  input_id: uint;
+  output_id: uint;
+  flags: uint;
+}
+
 table XNNStaticSlice {
   num_dims:uint;
   offsets:[uint];