Preserve logical element order when lowering view_copy

backends/cortex_m/test/build_test_runner.sh

@@ -48,6 +48,7 @@ ops_list=(
     aten::convolution.out
     aten::max_pool2d_with_indices.out
     dim_order_ops::_clone_dim_order.out
+    aten::view_copy.out
     aten::cat.out
     aten::full.out
     aten::ge.Tensor_out

backends/cortex_m/test/targets.bzl

@@ -135,15 +135,14 @@ def define_common_targets(is_fbcode = False):
                 "//executorch/backends/cortex_m/passes:cortex_passes",
                 "//executorch/backends/cortex_m/quantizer:quantizer",
                 "//executorch/exir:lib",
+                "//executorch/exir:memory",
                 "//executorch/exir:pass_base",
                 "//executorch/exir/_serialize:lib",
                 "//executorch/exir/dialects:lib",
                 "fbsource//third-party/pypi/pytest:pytest",
             ],
         )
 
-
-
         python_pytest(
             name = "test_fuse_conv_padding",
             srcs = ["test_fuse_conv_padding.py"],

backends/cortex_m/test/test_pass_manager.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.
@@ -12,7 +13,7 @@
 from executorch.backends.cortex_m.passes.cortex_m_pass_manager import CortexMPassManager
 from executorch.backends.cortex_m.quantizer.quantizer import CortexMQuantizer
 from executorch.backends.cortex_m.target_config import CortexM, CortexMTargetConfig
-from executorch.exir import to_edge, to_edge_transform_and_lower
+from executorch.exir import memory, to_edge, to_edge_transform_and_lower
 from executorch.exir._serialize._program import deserialize_pte_binary
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass
@@ -166,3 +167,50 @@ def call_operator(self, op, args, kwargs, meta):
     assert any(node.op == "get_attr" for node in program.graph.nodes)
     assert not program.graph_signature.buffers
     torch.testing.assert_close(program.module()(*inputs), torch.tensor([2.0, 4.0]))
+
+
+@pytest.mark.parametrize(
+    "channels_last,use_explicit_layout",
+    [(False, False), (True, False), (False, True)],
+)
+def test_flatten_preserves_patch_tokens(channels_last, use_explicit_layout):
+    class PatchTokens(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.conv = torch.nn.Conv2d(3, 4, 2, stride=2)
+
+        def forward(self, x):
+            return self.conv(x).flatten(2).transpose(1, 2)
+
+    torch.manual_seed(0)
+    model = PatchTokens().eval()
+    inputs = (torch.randn(1, 3, 6, 6),)
+    if channels_last:
+        model = model.to(memory_format=torch.channels_last)
+        inputs = (inputs[0].to(memory_format=torch.channels_last),)
+    exported, expected, scale = _quantize(model, inputs, use_explicit_layout)
+    edge = to_edge_transform_and_lower(
+        exported,
+        compile_config=cortex_m_edge_compile_config(),
+        transform_passes=CortexMPassManager(use_explicit_layout=use_explicit_layout),
+    )
+    torch.testing.assert_close(
+        edge.exported_program().module()(*inputs), expected, rtol=0, atol=2 * scale
+    )
+    program = edge.to_executorch().exported_program()
+    views = [node for node in program.graph.nodes if node.target == memory.view]
+    view_copies = [
+        node
+        for node in program.graph.nodes
+        if node.target == torch.ops.aten.view_copy.out
+    ]
+    assert views or view_copies
+    for node in views:
+        assert (
+            node.args[0].meta["val"].is_contiguous()
+        ), "Runtime views must not alias channels-last storage as NCHW"
+    if use_explicit_layout:
+        assert not any(
+            node.target == torch.ops.dim_order_ops._to_dim_order_copy.out
+            for node in program.graph.nodes
+        )

exir/passes/replace_view_copy_with_view_pass.py

@@ -434,6 +434,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
 
         This should be run after the NormalizeViewCopyBasePass.
 
+        Only contiguous inputs can share storage with the contiguous view spec.
         During memory planning, view nodes share the same storage as their base.
         """
 
@@ -449,6 +450,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
                 if (
                     _is_view_copy(node)
                     and all(u.op != "output" for u in node.users)
+                    and node.args[0].meta["val"].is_contiguous()
                     and is_copy_to_view_safe(node)
                 ):
                     base, _ = node.args
@@ -495,6 +497,7 @@ def ensures(self, graph_module: torch.fx.GraphModule) -> None:
                 assert not (
                     _is_view_copy(node)
                     and all(u.op != "output" for u in node.users)
+                    and node.args[0].meta["val"].is_contiguous()
                     and is_copy_to_view_safe(node)
                 )
                 if node.op == "call_function" and node.target == _VIEW_OP:

exir/tests/targets.bzl

@@ -1,3 +1,8 @@
+# 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.
+
 load("@fbcode_macros//build_defs:cpp_library.bzl", "cpp_library")
 load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
 load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest")
@@ -474,6 +479,7 @@ def define_common_targets(is_fbcode = False):
             "//executorch/exir/capture:config",
             "//executorch/exir/dialects:lib",
             "//executorch/exir/passes:lib",
+            "//executorch/extension/pybindings:portable_lib",  # @manual
         ],
     )
 

exir/tests/test_remove_view_copy.py

@@ -1,15 +1,17 @@
 # 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 copy
 import unittest
+from itertools import product
 
 import torch
 import torch.nn as nn
-from executorch.exir import memory, to_edge
+from executorch.exir import EdgeCompileConfig, memory, to_edge
 from executorch.exir.capture._config import ExecutorchBackendConfig
 from executorch.exir.dialects._ops import ops
 from executorch.exir.passes import MemoryPlanningPass, ToOutVarPass
@@ -22,6 +24,10 @@
 )
 from executorch.exir.passes.spec_prop_pass import SpecPropPass
 
+from executorch.extension.pybindings.portable_lib import (  # @manual
+    _load_for_executorch_from_buffer,
+)
+
 
 class TestModel1(nn.Module):
     __test__ = False
@@ -63,6 +69,78 @@ def _run_view_and_reinplace_passes(
         ReplaceViewCopyWithViewPass()(graph_module)
         return graph_module
 
+    def test_flatten_layout_runtime(self) -> None:
+        class Flatten(nn.Module):
+            def __init__(self, internal_channels_last):
+                super().__init__()
+                self.internal_channels_last = internal_channels_last
+
+            def forward(self, x):
+                if self.internal_channels_last:
+                    x = x.clone(memory_format=torch.channels_last)
+                return x.flatten(2).transpose(1, 2)
+
+        for layout, skip_dim_order, remove_view_copy in product(
+            ("contiguous", "channels_last", "singleton", "internal"),
+            (False, True),
+            (False, True),
+        ):
+            # Explicit layout conversions require dim-order operators.
+            if layout == "internal" and skip_dim_order:
+                continue
+            with self.subTest(
+                layout=layout,
+                skip_dim_order=skip_dim_order,
+                remove_view_copy=remove_view_copy,
+            ):
+                channels = 1 if layout == "singleton" else 3
+                x = torch.arange(16 * channels, dtype=torch.float32).reshape(
+                    2, channels, 2, 4
+                )
+                if layout in ("channels_last", "singleton"):
+                    x = x.contiguous(memory_format=torch.channels_last)
+                model = Flatten(layout == "internal").eval()
+                edge = to_edge(
+                    torch.export.export(model, (x,), strict=True),
+                    compile_config=EdgeCompileConfig(_skip_dim_order=skip_dim_order),
+                )
+                program = edge.to_executorch(
+                    config=ExecutorchBackendConfig(remove_view_copy=remove_view_copy)
+                )
+                targets = [
+                    node.target for node in program.exported_program().graph.nodes
+                ]
+                uses_view = remove_view_copy and layout in ("contiguous", "singleton")
+                self.assertEqual(memory.view in targets, uses_view)
+                self.assertEqual(torch.ops.aten.view_copy.out in targets, not uses_view)
+                if skip_dim_order and layout != "internal":
+                    self.assertFalse(any("dim_order" in str(t) for t in targets))
+                    self.assertNotIn(torch.ops.aten._to_copy.out, targets)
+                runtime = _load_for_executorch_from_buffer(program.buffer)
+                actual = runtime.run_method("forward", (x,))[0]
+                torch.testing.assert_close(actual, model(x), rtol=0, atol=0)
+
+    def test_view_replacement_preserves_input_layout(self) -> None:
+        class Flatten(nn.Module):
+            def forward(self, x):
+                return x.flatten(2).transpose(1, 2)
+
+        for channels in (1, 3):
+            for channels_last in (False, True):
+                with self.subTest(channels=channels, channels_last=channels_last):
+                    x = torch.arange(8 * channels, dtype=torch.float32).reshape(
+                        1, channels, 2, 4
+                    )
+                    if channels_last:
+                        x = x.contiguous(memory_format=torch.channels_last)
+                    graph = self._run_view_and_reinplace_passes(Flatten(), (x,))
+                    targets = [node.target for node in graph.graph.nodes]
+                    if channels_last and channels > 1:
+                        self.assertNotIn(memory.view, targets)
+                        self.assertIn(ops.edge.aten.view_copy.default, targets)
+                    else:
+                        self.assertIn(memory.view, targets)
+
     def test_disable(self) -> None:
         model = TestModel1()
         model.eval()

kernels/portable/cpu/op_view_copy.cpp

@@ -1,6 +1,7 @@
 /*
  * 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,6 +10,7 @@
 #include <cstdint>
 #include <cstring>
 
+#include <executorch/kernels/portable/cpu/util/broadcast_indexes_range.h>
 #include <executorch/kernels/portable/cpu/util/copy_ops_util.h>
 #include <executorch/runtime/kernel/kernel_includes.h>
 
@@ -44,16 +46,29 @@ Tensor& view_copy_out(
       out,
       "Failed to resize output tensor.");
 
-  ET_KERNEL_CHECK(
-      ctx, tensors_have_same_dim_order(self, out), InvalidArgument, out);
-
-  ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(self), InvalidArgument, out);
+  ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(out), InvalidArgument, out);
 
   ET_KERNEL_CHECK(
       ctx, check_view_copy_args(self, size_int64_t, out), InvalidArgument, out);
 
-  if (self.nbytes() > 0) {
+  if (self.nbytes() == 0) {
+    return out;
+  }
+  if (tensor_is_default_dim_order(self)) {
     memcpy(out.mutable_data_ptr(), self.const_data_ptr(), self.nbytes());
+  } else {
+    const auto* input_data = self.const_data_ptr<uint8_t>();
+    auto* output_data = out.mutable_data_ptr<uint8_t>();
+    const auto element_size = self.element_size();
+    // Iterate the input's logical shape; the output may have a different rank.
+    for (const auto [output_index, input_index] :
+         BroadcastIndexesRange<1, /*support_noncontiguous_input_tensors=*/true>(
+             self, self)) {
+      memcpy(
+          output_data + output_index * element_size,
+          input_data + input_index * element_size,
+          element_size);
+    }
   }
   return out;
 }

kernels/test/op_view_copy_test.cpp

@@ -1,6 +1,7 @@
 /*
  * 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.
@@ -146,6 +147,39 @@ TEST_F(OpViewTest, EmptyInputSupported) {
 #undef TEST_ENTRY
 }
 
+TEST_F(OpViewTest, ChannelsLastInputPreservesLogicalOrder) {
+  TensorFactory<ScalarType::Float> tf;
+  for (const auto& shape :
+       std::vector<std::vector<int32_t>>{{1, 2, 1, 3}, {1, 2, 1, 1, 3}}) {
+    Tensor input = tf.make_channels_last(shape, {0, 3, 1, 4, 2, 5});
+    Tensor out = tf.zeros({2, 3});
+    Tensor expected = tf.make({2, 3}, {0, 1, 2, 3, 4, 5});
+    const int64_t size[] = {-1, 3};
+    op_view_copy_out(input, size, out);
+    EXPECT_TENSOR_EQ(out, expected);
+  }
+}
+
+TEST_F(OpViewTest, EmptyChannelsLastInput) {
+  TensorFactory<ScalarType::Float> tf;
+  Tensor input = tf.make_channels_last({1, 2, 0, 3}, {});
+  Tensor out = tf.zeros({2, 0});
+  const int64_t size[] = {2, 0};
+  op_view_copy_out(input, size, out);
+  EXPECT_TENSOR_EQ(out, tf.zeros({2, 0}));
+}
+
+TEST_F(OpViewTest, ChannelsLastInputResizesOutput) {
+  TensorFactory<ScalarType::Int> tf;
+  Tensor input = tf.make_channels_last({1, 2, 1, 3}, {0, 3, 1, 4, 2, 5});
+  Tensor out =
+      tf.zeros({4, 4}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
+  Tensor expected = tf.make({2, 3}, {0, 1, 2, 3, 4, 5});
+  const int64_t size[] = {2, 3};
+  op_view_copy_out(input, size, out);
+  EXPECT_TENSOR_EQ(out, expected);
+}
+
 TEST_F(OpViewTest, InputOutputMismatchedSizesDie) {
   TensorFactory<ScalarType::Int> tf;
   std::vector<int32_t> size_in = {3, 1, 1, 2};

shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl

@@ -1,3 +1,8 @@
+# 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.
+
 load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "get_aten_mode_options", "is_xplat", "runtime")
 load("@fbsource//xplat/executorch/build:selects.bzl", "selects")
 
@@ -1394,6 +1399,7 @@ ATEN_OPS = (
     op_target(
         name = "op_view_copy",
         deps = [
+            "//executorch/kernels/portable/cpu/util:broadcast_indexes_range",
             "//executorch/kernels/portable/cpu/util:copy_ops_util",
         ],
     ),