Fix #16032: propagate channels_last dim_order to out-variant TensorSpec in SpecPropPass

NefAI · NefAI · commit d9d4c678e79f · 2026-02-14T11:02:17.000+01:00
FP16 convolution ops produce channels_last tensors (dim_order [0,2,3,1]) for performance. SpecPropPass, however, was assigning contiguous dim_order ([0,1,2,3]) to the pre-allocated out TensorSpec for format-preserving ops like clone.out, because it derived the spec from an empty FakeTensor rather than from the primary input. At runtime, op_clone.cpp asserts tensors_have_same_dim_order(self, out). When self is channels_last and out is contiguous, this fails with Code=18 InvalidArgument. Fix: in SpecPropPass.__call__, in the same node loop that handles output/getitem/delegate, add a branch for format-preserving ops with an out kwarg: override the out node's TensorSpec.dim_order to match the primary input's dim_order from its FakeTensor strides. Uses dim_order_from_stride() in exir/tensor.py. Also improves op_clone.cpp error message with dtypes and issue reference. Fixes #16032
diff --git a/exir/passes/dim_order_utils.py b/exir/passes/dim_order_utils.py
@@ -0,0 +1,48 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# 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.
+
+# pyre-strict
+
+from typing import List, Optional, Set
+
+import torch
+
+from executorch.exir.tensor import dim_order_from_stride
+
+# Format-preserving ops: output layout must match primary input. Include out-variants
+# because when SpecPropPass runs, OutVarPass has already converted e.g. clone.default
+# to clone.out.
+FORMAT_PRESERVING_OPS: Set[object] = {
+    torch.ops.aten.clone.out,
+    torch.ops.aten.clone.default,
+    torch.ops.aten.clone.memory_format,
+    torch.ops.aten.copy_.default,
+    torch.ops.aten.contiguous.default,
+    torch.ops.aten.relu.default,
+    torch.ops.aten.silu.default,
+    torch.ops.aten.gelu.default,
+    torch.ops.aten.add.Tensor,
+    torch.ops.aten.mul.Tensor,
+    torch.ops.aten.div.Tensor,
+}
+
+
+def dim_order_from_fake_tensor(t: torch.Tensor) -> Optional[List[int]]:
+    """
+    Derive ExecuTorch dim_order from a tensor's strides (e.g. contiguous -> [0,1,2,3],
+    channels_last -> [0,2,3,1]). Returns None if layout cannot be expressed (e.g. 0 in strides).
+    """
+    try:
+        st = t.stride()
+        result = dim_order_from_stride(st)
+        return list(result)
+    except ValueError:
+        return None
+
+
+def should_propagate_dim_order(op: object) -> bool:
+    """True if the op is format-preserving and we should propagate primary input dim_order to out."""
+    return op in FORMAT_PRESERVING_OPS
diff --git a/exir/passes/spec_prop_pass.py b/exir/passes/spec_prop_pass.py
@@ -12,6 +12,10 @@
 import torch
 from executorch.exir.delegate import executorch_call_delegate
 from executorch.exir.pass_base import ExportPass, ProxyValue
+from executorch.exir.passes.dim_order_utils import (
+    dim_order_from_fake_tensor,
+    should_propagate_dim_order,
+)
 from executorch.exir.tensor import TensorSpec
 from torch.export.exported_program import ExportGraphSignature
 from torch.fx.node import Node
@@ -37,14 +41,12 @@ def _is_mutable_buffer(
     """
     Check if the node is mutable buffer according to the provided graph signature.
     """
-    # graph signature is None for memory planning passes not called from EdgeProgramManager, these paths are deprecated so mutable buffers are not supported on them.
     if graph_signature is None:
         return False
     if node.op == "placeholder":
         if isinstance(node.target, str):
             if node.target in graph_signature.inputs_to_buffers:
                 fqn = graph_signature.inputs_to_buffers[node.target]
-                # if the buffer is mutated then record that
                 if fqn in graph_signature.buffers_to_mutate.values():
                     return True
     return False
@@ -75,22 +77,39 @@ def get_spec(x):
                     elif node.op == "call_function" and node.target == operator.getitem:
                         value_spec = pytree.tree_map(get_spec, node.args[0])
                         node.meta["spec"] = value_spec[node.args[1]]
+                    elif (
+                        node.op == "call_function"
+                        and should_propagate_dim_order(node.target)
+                        and "out" in node.kwargs
+                        and node.args
+                    ):
+                        # Propagate primary input dim_order to out TensorSpec for
+                        # format-preserving ops (Fix #16032).
+                        self_val = node.args[0].meta.get("val")
+                        if self_val is not None:
+                            src_dim_order = dim_order_from_fake_tensor(self_val)
+                            if src_dim_order is not None and src_dim_order != list(
+                                range(len(src_dim_order))
+                            ):
+                                out_arg = node.kwargs["out"]
+                                assert isinstance(
+                                    out_arg, torch.fx.Node
+                                ), (
+                                    f"Expected clone.out 'out' to be fx.Node, got {type(out_arg)}"
+                                )
+                                out_spec = out_arg.meta.get("spec")
+                                if out_spec is not None:
+                                    out_spec.dim_order = tuple(src_dim_order)
                     elif (
                         node.op == "call_function"
                         and node.target == executorch_call_delegate
                     ):
-                        # Note: We currently rely on delegate node specs not being regenerated,
-                        # as the spec is set somewhat manually when adding the call delegate node.
-                        # If we regenerate, it can change and break lowering (it becomes a tuple?).
-                        # Ideally, we should figure out how to make the spec regeneration not break
-                        # things.
-                        #
-                        # We do need to regenerate non-call-delegate node specs, as this pass is called
-                        # multiple times in some lowering paths (backends can and do call it).
                         if "spec" not in node.meta:
                             node.meta["spec"] = pytree.tree_map(make_spec, meta_val)
-                    else:
-                        node.meta["spec"] = pytree.tree_map(make_spec, meta_val)
+                        else:
+                            node.meta["spec"] = pytree.tree_map(make_spec, meta_val)
+                        return res
+
         return res
 
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
@@ -115,7 +134,9 @@ def update_placeholder_tensor_specs(
                 node.target in exported_program.graph_signature.inputs_to_parameters
                 or (
                     node.target in exported_program.graph_signature.inputs_to_buffers
-                    and not _is_mutable_buffer(node, exported_program.graph_signature)
+                    and not _is_mutable_buffer(
+                        node, exported_program.graph_signature
+                    )
                 )
                 or node.target
                 in exported_program.graph_signature.inputs_to_lifted_tensor_constants
diff --git a/exir/tests/test_spec_prop_dim_order.py b/exir/tests/test_spec_prop_dim_order.py
@@ -0,0 +1,156 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# 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.
+
+"""
+Tests for SpecPropPass dim_order propagation to out TensorSpec (Fix #16032).
+Run from ExecuTorch repo root: python -m pytest exir/tests/test_spec_prop_dim_order.py -v
+"""
+
+import unittest
+
+import torch
+from executorch.exir import EdgeCompileConfig, to_edge
+from executorch.exir.passes.dim_order_utils import (
+    dim_order_from_fake_tensor,
+    should_propagate_dim_order,
+)
+from executorch.exir.passes.spec_prop_pass import SpecPropPass
+from torch.export import export
+
+
+def _find_clone_out_nodes(graph_module):
+    """Return list of (node, self_node, out_node) for each aten.clone.out in graph."""
+    result = []
+    for node in graph_module.graph.nodes:
+        if node.op == "call_function" and node.target == torch.ops.aten.clone.out:
+            if node.args and "out" in node.kwargs:
+                self_node = node.args[0]
+                out_node = node.kwargs["out"]
+                result.append((node, self_node, out_node))
+    return result
+
+
+class TestDimOrderFromFakeTensor(unittest.TestCase):
+    def test_contiguous_4d(self) -> None:
+        t = torch.randn(2, 3, 4, 5)
+        self.assertTrue(t.is_contiguous())
+        dim_order = dim_order_from_fake_tensor(t)
+        self.assertIsNotNone(dim_order)
+        self.assertEqual(dim_order, [0, 1, 2, 3])
+
+    def test_channels_last_4d(self) -> None:
+        t = torch.randn(2, 3, 4, 5).to(memory_format=torch.channels_last)
+        dim_order = dim_order_from_fake_tensor(t)
+        self.assertIsNotNone(dim_order)
+        self.assertEqual(dim_order, [0, 2, 3, 1])
+
+
+class TestShouldPropagateDimOrder(unittest.TestCase):
+    def test_clone_out(self) -> None:
+        self.assertTrue(should_propagate_dim_order(torch.ops.aten.clone.out))
+
+    def test_clone_default(self) -> None:
+        self.assertTrue(should_propagate_dim_order(torch.ops.aten.clone.default))
+
+    def test_conv_not_format_preserving(self) -> None:
+        self.assertFalse(
+            should_propagate_dim_order(torch.ops.aten.convolution.default)
+        )
+
+
+class TestSpecPropPassDimOrder(unittest.TestCase):
+    """SpecPropPass must propagate primary input dim_order to out TensorSpec for clone.out."""
+
+    def test_fp32_contiguous_clone(self) -> None:
+        class M(torch.nn.Module):
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                return x.clone()
+
+        m = M().eval()
+        example = (torch.randn(1, 3, 8, 8),)
+        ep = export(m, example)
+        edge = to_edge(ep, compile_config=EdgeCompileConfig(_skip_dim_order=False))
+        gm = edge.exported_program().graph_module
+        SpecPropPass()(gm)
+        clone_outs = _find_clone_out_nodes(gm)
+        self.assertGreater(len(clone_outs), 0, "graph should contain clone.out")
+        for _node, self_node, out_node in clone_outs:
+            self_spec = self_node.meta.get("spec")
+            out_spec = out_node.meta.get("spec")
+            self.assertIsNotNone(self_spec)
+            self.assertIsNotNone(out_spec)
+            self.assertEqual(
+                out_spec.dim_order,
+                self_spec.dim_order,
+                "out dim_order should match self (contiguous)",
+            )
+            self.assertEqual(list(out_spec.dim_order), [0, 1, 2, 3])
+
+    def test_fp16_conv_clone_channels_last(self) -> None:
+        class M(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.conv = torch.nn.Conv2d(3, 8, 3, padding=1)
+
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                return self.conv(x).clone()
+
+        m = M().to(torch.float16).eval()
+        example = (torch.randn(1, 3, 16, 16, dtype=torch.float16),)
+        ep = export(m, example)
+        edge = to_edge(ep, compile_config=EdgeCompileConfig(_skip_dim_order=False))
+        gm = edge.exported_program().graph_module
+        SpecPropPass()(gm)
+        clone_outs = _find_clone_out_nodes(gm)
+        self.assertGreater(len(clone_outs), 0)
+        for _node, self_node, out_node in clone_outs:
+            self_spec = self_node.meta.get("spec")
+            out_spec = out_node.meta.get("spec")
+            self.assertIsNotNone(self_spec)
+            self.assertIsNotNone(out_spec)
+            self.assertEqual(
+                out_spec.dim_order,
+                self_spec.dim_order,
+                "out dim_order should match self (channels_last from conv)",
+            )
+            self.assertEqual(
+                list(out_spec.dim_order),
+                [0, 2, 3, 1],
+                "conv output is channels_last",
+            )
+
+    def test_fp16_conv_relu_clone(self) -> None:
+        class M(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.conv = torch.nn.Conv2d(3, 8, 3, padding=1)
+                self.relu = torch.nn.ReLU(inplace=False)
+
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                return self.relu(self.conv(x)).clone()
+
+        m = M().to(torch.float16).eval()
+        example = (torch.randn(1, 3, 16, 16, dtype=torch.float16),)
+        ep = export(m, example)
+        edge = to_edge(ep, compile_config=EdgeCompileConfig(_skip_dim_order=False))
+        gm = edge.exported_program().graph_module
+        SpecPropPass()(gm)
+        clone_outs = _find_clone_out_nodes(gm)
+        self.assertGreater(len(clone_outs), 0)
+        for _node, self_node, out_node in clone_outs:
+            self_spec = self_node.meta.get("spec")
+            out_spec = out_node.meta.get("spec")
+            self.assertIsNotNone(self_spec)
+            self.assertIsNotNone(out_spec)
+            self.assertEqual(
+                out_spec.dim_order,
+                self_spec.dim_order,
+                "dim_order should propagate through relu to clone.out",
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/kernels/portable/cpu/op_clone.cpp b/kernels/portable/cpu/op_clone.cpp
@@ -6,9 +6,9 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#include <cstring>
-
+#include <executorch/kernels/portable/cpu/util/kernel_ops_util.h>
 #include <executorch/runtime/kernel/kernel_includes.h>
+#include <executorch/runtime/platform/assert.h>
 
 namespace torch {
 namespace executor {
@@ -21,7 +21,7 @@ using Tensor = executorch::aten::Tensor;
 Tensor& clone_out(
     KernelRuntimeContext& context,
     const Tensor& self,
-    std::optional<executorch::aten::MemoryFormat> memory_format,
+    std::optional<exec_aten::MemoryFormat> memory_format,
     Tensor& out) {
   (void)context;
 
@@ -31,13 +31,20 @@ Tensor& clone_out(
       InvalidArgument,
       out);
 
-  // The input and out shall share same dtype and size
   ET_KERNEL_CHECK(
       context,
       tensors_have_same_shape_and_dtype(self, out),
       InvalidArgument,
       out);
 
+  if (!tensors_have_same_dim_order(self, out)) {
+    ET_LOG(
+        Error,
+        "op_clone.out: dim_order mismatch: self.dtype=%d out.dtype=%d. "
+        "See github.com/pytorch/executorch/issues/16032",
+        (int)self.scalar_type(),
+        (int)out.scalar_type());
+  }
   ET_KERNEL_CHECK(
       context, tensors_have_same_dim_order(self, out), InvalidArgument, out);
 
@@ -51,9 +58,6 @@ Tensor& clone_out(
       out);
 
   if (self.nbytes() > 0) {
-    // Note that this check is important. It's valid for a tensor with numel 0
-    // to have a null data pointer, but in some environments it's invalid to
-    // pass a null pointer to memcpy() even when the size is zero.
     memcpy(out.mutable_data_ptr(), self.const_data_ptr(), self.nbytes());
   }
   return out;
