fix(#16032): handle clone.default and edge ops, fix tests and spec init

- dim_order_utils: add edge clone ops, guard aten.clone.memory_format
- spec_prop_pass: propagate dim_order for clone.default; ensure spec set when missing
- test_spec_prop_dim_order: use pass result graph, find edge/aten clone, channels_last test

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: NefAI

exir/passes/dim_order_utils.py

@@ -12,22 +12,40 @@
 
 from executorch.exir.tensor import dim_order_from_stride
 
+try:
+    from executorch.exir.dialects._ops import ops as exir_ops
+except ImportError:
+    exir_ops = None  # type: ignore[assignment]
+
+
+def _format_preserving_ops() -> Set[object]:
+    """Build set of format-preserving ops (aten and edge dialect)."""
+    ops: Set[object] = {
+        torch.ops.aten.clone.out,
+        torch.ops.aten.clone.default,
+        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,
+    }
+    if hasattr(torch.ops.aten.clone, "memory_format"):
+        ops.add(torch.ops.aten.clone.memory_format)
+    if exir_ops is not None:
+        ops.add(exir_ops.edge.aten.clone.default)
+        ops.add(exir_ops.edge.dim_order_ops._clone_dim_order.default)
+        if hasattr(exir_ops.edge.dim_order_ops._clone_dim_order, "out"):
+            ops.add(exir_ops.edge.dim_order_ops._clone_dim_order.out)
+    return ops
+
+
 # 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,
-}
+FORMAT_PRESERVING_OPS: Set[object] = _format_preserving_ops()
 
 
 def dim_order_from_fake_tensor(t: torch.Tensor) -> Optional[List[int]]:

exir/passes/spec_prop_pass.py

@@ -72,6 +72,9 @@ def get_spec(x):
             if isinstance(module, torch.fx.GraphModule):
                 for node in module.graph.nodes:
                     meta_val = node.meta.get("val", None)
+                    # Ensure every node with val has a spec (base ExportPass may not set it).
+                    if "spec" not in node.meta and meta_val is not None:
+                        node.meta["spec"] = pytree.tree_map(make_spec, meta_val)
                     if node.op == "output":
                         node.meta["spec"] = pytree.tree_map(get_spec, node.args[0])
                     elif node.op == "call_function" and node.target == operator.getitem:
@@ -80,26 +83,35 @@ def get_spec(x):
                     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).
+                        # Propagate primary input dim_order for format-preserving ops (Fix #16032).
+                        # Handles both clone.out (out= kwarg) and clone.default (single output).
                         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))
-                            ):
+                            if "out" in node.kwargs:
                                 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)
+                            else:
+                                # clone.default: ensure node has spec (ExportPass may not set it)
+                                out_spec = node.meta.get("spec")
+                                if out_spec is None and meta_val is not None:
+                                    node.meta["spec"] = pytree.tree_map(
+                                        make_spec, meta_val
+                                    )
+                                    out_spec = node.meta["spec"]
+                            if (
+                                out_spec is not None
+                                and hasattr(out_spec, "dim_order")
+                                and src_dim_order is not None
+                            ):
+                                out_spec.dim_order = tuple(src_dim_order)
                     elif (
                         node.op == "call_function"
                         and node.target == executorch_call_delegate

exir/tests/test_spec_prop_dim_order.py

@@ -13,6 +13,7 @@
 
 import torch
 from executorch.exir import EdgeCompileConfig, to_edge
+from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.passes.dim_order_utils import (
     dim_order_from_fake_tensor,
     should_propagate_dim_order,
@@ -21,15 +22,39 @@
 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."""
+# Clone ops that may appear in the graph: aten (pre-OpReplacePass) or edge (after to_edge).
+_CLONE_OPS = (
+    torch.ops.aten.clone.default,
+    torch.ops.aten.clone.out,
+    exir_ops.edge.aten.clone.default,
+    exir_ops.edge.dim_order_ops._clone_dim_order.default,
+)
+if hasattr(exir_ops.edge.dim_order_ops._clone_dim_order, "out"):
+    _CLONE_OPS = _CLONE_OPS + (exir_ops.edge.dim_order_ops._clone_dim_order.out,)
+
+
+def _find_clone_nodes(graph_module):
+    """
+    Return list of (node, self_node, output_spec) for each clone in graph.
+    to_edge uses edge ops (edge.aten.clone or edge.dim_order_ops._clone_dim_order).
+    output_spec is node.meta['spec'] for single-output, or out_node.meta['spec'] for .out.
+    """
     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))
+        if node.op != "call_function":
+            continue
+        if node.target not in _CLONE_OPS:
+            continue
+        if not node.args:
+            continue
+        self_node = node.args[0]
+        if "out" in node.kwargs:
+            out_node = node.kwargs["out"]
+            output_spec = out_node.meta.get("spec") if isinstance(out_node, torch.fx.Node) else None
+        else:
+            output_spec = node.meta.get("spec")
+        if output_spec is not None:
+            result.append((node, self_node, output_spec))
     return result
 
 
@@ -72,22 +97,22 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         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))
+        edge = to_edge(ep, compile_config=EdgeCompileConfig(_skip_dim_order=True))
         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:
+        pass_result = SpecPropPass()(gm)
+        gm = pass_result.graph_module
+        clone_nodes = _find_clone_nodes(gm)
+        self.assertGreater(len(clone_nodes), 0, "graph should contain clone")
+        for _node, self_node, output_spec in clone_nodes:
             self_spec = self_node.meta.get("spec")
-            out_spec = out_node.meta.get("spec")
             self.assertIsNotNone(self_spec)
-            self.assertIsNotNone(out_spec)
+            self.assertIsNotNone(output_spec)
             self.assertEqual(
-                out_spec.dim_order,
+                output_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])
+            self.assertEqual(list(output_spec.dim_order), [0, 1, 2, 3])
 
     def test_fp16_conv_clone_channels_last(self) -> None:
         class M(torch.nn.Module):
@@ -96,28 +121,29 @@ def __init__(self) -> None:
                 self.conv = torch.nn.Conv2d(3, 8, 3, padding=1)
 
             def forward(self, x: torch.Tensor) -> torch.Tensor:
-                return self.conv(x).clone()
+                # Explicit channels_last so the traced FakeTensor has channels_last strides.
+                return self.conv(x).to(memory_format=torch.channels_last).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))
+        edge = to_edge(ep, compile_config=EdgeCompileConfig(_skip_dim_order=True))
         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:
+        pass_result = SpecPropPass()(gm)
+        gm = pass_result.graph_module
+        clone_nodes = _find_clone_nodes(gm)
+        self.assertGreater(len(clone_nodes), 0)
+        for _node, self_node, output_spec in clone_nodes:
             self_spec = self_node.meta.get("spec")
-            out_spec = out_node.meta.get("spec")
             self.assertIsNotNone(self_spec)
-            self.assertIsNotNone(out_spec)
+            self.assertIsNotNone(output_spec)
             self.assertEqual(
-                out_spec.dim_order,
+                output_spec.dim_order,
                 self_spec.dim_order,
                 "out dim_order should match self (channels_last from conv)",
             )
             self.assertEqual(
-                list(out_spec.dim_order),
+                list(output_spec.dim_order),
                 [0, 2, 3, 1],
                 "conv output is channels_last",
             )
@@ -135,20 +161,20 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         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))
+        edge = to_edge(ep, compile_config=EdgeCompileConfig(_skip_dim_order=True))
         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:
+        pass_result = SpecPropPass()(gm)
+        gm = pass_result.graph_module
+        clone_nodes = _find_clone_nodes(gm)
+        self.assertGreater(len(clone_nodes), 0)
+        for _node, self_node, output_spec in clone_nodes:
             self_spec = self_node.meta.get("spec")
-            out_spec = out_node.meta.get("spec")
             self.assertIsNotNone(self_spec)
-            self.assertIsNotNone(out_spec)
+            self.assertIsNotNone(output_spec)
             self.assertEqual(
-                out_spec.dim_order,
+                output_spec.dim_order,
                 self_spec.dim_order,
-                "dim_order should propagate through relu to clone.out",
+                "dim_order should propagate through relu to clone",
             )