Fix #16032: propagate channels_last dim_order to out-variant TensorSp…

exir/_serialize/_flatbuffer.py

@@ -197,6 +197,9 @@ def _run_flatc(args: Sequence[str]) -> None:
     else:
         # Expect the `flatc` tool to be on the system path or set as an env var.
         flatc_path = os.getenv("FLATC_EXECUTABLE")
+        if flatc_path and not os.path.isfile(flatc_path):
+            # Env set to a path that doesn't exist (e.g. placeholder); use PATH.
+            flatc_path = "flatc"
         if not flatc_path:
             flatc_path = "flatc"
         subprocess.run([flatc_path] + list(args), check=True)

exir/backend/test/test_compatibility.py

@@ -19,12 +19,25 @@
 )
 
 from executorch.extension.pybindings.portable_lib import (
+    _get_registered_backend_names,  # @manual
     _load_for_executorch_from_buffer,  # @manual
 )
 from torch.export import export
 
 
+def _has_backend_with_compiler_demo() -> bool:
+    """Check if BackendWithCompilerDemo is linked into the portable runtime."""
+    try:
+        return "BackendWithCompilerDemo" in _get_registered_backend_names()
+    except Exception:
+        return False
+
+
 class TestCompatibility(unittest.TestCase):
+    @unittest.skipUnless(
+        _has_backend_with_compiler_demo(),
+        "BackendWithCompilerDemo not registered (build with EXECUTORCH_BUILD_TESTS=ON)",
+    )
     def test_compatibility_in_runtime(self):
         class SinModule(torch.nn.Module):
             def __init__(self):
@@ -70,6 +83,10 @@ def forward(self, x):
         ):
             executorch_module.run_method("forward")
 
+    @unittest.skipUnless(
+        _has_backend_with_compiler_demo(),
+        "BackendWithCompilerDemo not registered (build with EXECUTORCH_BUILD_TESTS=ON)",
+    )
     def test_compatibility_in_runtime_edge_program_manager(self):
         class SinModule(torch.nn.Module):
             def __init__(self):

exir/backend/test/test_lowered_backend_module.py

@@ -20,11 +20,20 @@
 from executorch.exir.schema import DelegateCall, Program
 
 from executorch.extension.pybindings.portable_lib import (  # @manual
+    _get_registered_backend_names,
     _load_for_executorch_from_buffer,
 )
 from torch.export import export
 
 
+def _has_backend_with_compiler_demo() -> bool:
+    """Check if BackendWithCompilerDemo is linked into the portable runtime."""
+    try:
+        return "BackendWithCompilerDemo" in _get_registered_backend_names()
+    except Exception:
+        return False
+
+
 class TestBackendAPI(unittest.TestCase):
     def validate_lowered_module_program(self, program: Program) -> None:
         """
@@ -64,6 +73,10 @@ def forward(self, *args):
             .executorch_program
         )
 
+    @unittest.skipUnless(
+        _has_backend_with_compiler_demo(),
+        "BackendWithCompilerDemo not registered (build with EXECUTORCH_BUILD_TESTS=ON)",
+    )
     def test_emit_lowered_backend_module_end_to_end(self):
         class SinModule(torch.nn.Module):
             def __init__(self):

exir/backend/test/test_to_backend_multi_method.py

@@ -39,13 +39,22 @@
     Program,
 )
 from executorch.extension.pybindings.portable_lib import (  # @manual
+    _get_registered_backend_names,
     _load_for_executorch_from_buffer,
 )
 from torch.export.exported_program import ExportedProgram
 
 from torch.testing import FileCheck
 
 
+def _has_backend_with_compiler_demo() -> bool:
+    """Check if BackendWithCompilerDemo is linked into the portable runtime."""
+    try:
+        return "BackendWithCompilerDemo" in _get_registered_backend_names()
+    except Exception:
+        return False
+
+
 class TestToBackendMultiMethod(unittest.TestCase):
     """
     Testing suite used to test multi method to_backend lowering. The test suite uses demo backends
@@ -504,6 +513,10 @@ def forward(self, x):
         ):
             self._test(test_set)
 
+    @unittest.skipUnless(
+        _has_backend_with_compiler_demo(),
+        "BackendWithCompilerDemo not registered (build with EXECUTORCH_BUILD_TESTS=ON)",
+    )
     def test_multi_method_end_to_end(self):
         """
         Tests multi method lowering end-to-end. Lowers the same Sin Module for two methods

exir/lowered_backend_module.py

@@ -288,6 +288,8 @@ def program(
                 args=(delegate_node, i),
             )
             getitem_node.meta["val"] = delegate_node.meta["val"][i]
+            # FIX: Set spec at creation so SpecPropPass/MemoryPlanningPass don't need to synthesize it (issue #16032).
+            getitem_node.meta["spec"] = make_spec(delegate_node.meta["val"][i])
             getitem_nodes.append(getitem_node)
         lowered_exported_program.graph.output(getitem_nodes)
 

exir/passes/replace_view_copy_with_view_pass.py

@@ -283,7 +283,6 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
                     # Create spec for the node.
                     # _ViewSpec gives a view into its base spec for non-size
                     # related information.
-
                     # the shape is not the same as node.args[1] because node.args[1]
                     # can have an inferred sizes (-1).
                     shape = node.meta["val"].shape

exir/passes/spec_prop_pass.py

@@ -12,13 +12,124 @@
 import torch
 from executorch.exir.delegate import executorch_call_delegate
 from executorch.exir.pass_base import ExportPass, ProxyValue
-from executorch.exir.tensor import TensorSpec
+from executorch.exir.tensor import TensorSpec, dim_order_from_stride, stride_from_dim_order
 from torch.export.exported_program import ExportGraphSignature
 from torch.fx.node import Node
 from torch.fx.passes.infra.pass_base import PassResult
 from torch.utils import _pytree as pytree
 
 
+# Ops that TRANSFORM layout — output dim_order from op's explicit dim_order kwarg.
+# Source: ExecuTorch issues #8037 and #6330. Verified (Q3): torch.ops.dim_order_ops.
+try:
+    _LAYOUT_TRANSFORMING_OPS = frozenset({
+        torch.ops.dim_order_ops._to_dim_order_copy.default,
+        torch.ops.dim_order_ops._clone_dim_order.default,
+    })
+    _LAYOUT_TRANSFORMING_OP_NAMES = frozenset({"dim_order_ops::_to_dim_order_copy", "dim_order_ops::_clone_dim_order"})
+except AttributeError:
+    _LAYOUT_TRANSFORMING_OPS = frozenset()
+    _LAYOUT_TRANSFORMING_OP_NAMES = frozenset()
+
+
+def _is_layout_transforming_op(target) -> bool:
+    """True if op is layout-transforming (by identity, schema name, or string)."""
+    if target in _LAYOUT_TRANSFORMING_OPS:
+        return True
+    try:
+        schema = getattr(target, "_schema", None)
+        if schema is not None and schema.name is not None:
+            if schema.name in _LAYOUT_TRANSFORMING_OP_NAMES:
+                return True
+    except Exception:
+        pass
+    s = str(target)
+    if "_to_dim_order_copy" in s or "_clone_dim_order" in s:
+        return True
+    return False
+
+# Ops where output memory format is IDENTICAL to primary input. Reference: PyTorch docs memory_format.
+_FORMAT_PRESERVING_OPS: frozenset = frozenset({
+    torch.ops.aten.clone.default,
+    torch.ops.aten.clone.out,
+    torch.ops.aten.relu.default,
+    torch.ops.aten.relu.out,
+    torch.ops.aten.relu_.default,
+    torch.ops.aten.silu.default,
+    torch.ops.aten.silu.out,
+    torch.ops.aten.silu_.default,
+    torch.ops.aten.gelu.default,
+    torch.ops.aten.gelu.out,
+    torch.ops.aten.neg.default,
+    torch.ops.aten.neg.out,
+    torch.ops.aten.abs.default,
+    torch.ops.aten.abs.out,
+    torch.ops.aten.exp.default,
+    torch.ops.aten.exp.out,
+    torch.ops.aten.sqrt.default,
+    torch.ops.aten.sqrt.out,
+    torch.ops.aten.rsqrt.default,
+    torch.ops.aten.rsqrt.out,
+})
+
+assert _LAYOUT_TRANSFORMING_OPS.isdisjoint(_FORMAT_PRESERVING_OPS), (
+    "Op appears in both _LAYOUT_TRANSFORMING_OPS and _FORMAT_PRESERVING_OPS — check classification"
+)
+
+
+def _get_primary_tensor_input(node: Node) -> Optional[Node]:
+    """First argument that is an fx.Node with a FakeTensor val (primary input for layout)."""
+    for arg in node.args:
+        if (
+            isinstance(arg, Node)
+            and isinstance(arg.meta.get("val"), torch.Tensor)
+        ):
+            return arg
+    return None
+
+
+def _fix_out_spec_dim_order(node: Node) -> None:
+    """
+    For out-variant nodes, set the out kwarg node's TensorSpec.dim_order to the
+    layout the op will produce. For layout-transforming ops that return the
+    result (no out=), set this node's spec.dim_order from the dim_order kwarg.
+    Also updates spec.stride to be consistent with the new dim_order.
+    Fixes Code=18 at runtime (issue #16032).
+    """
+    # Layout-transforming ops: set this node's spec from dim_order kwarg (return-value case)
+    if _is_layout_transforming_op(node.target):
+        explicit_dim_order = node.kwargs.get("dim_order")
+        if explicit_dim_order is not None:
+            spec = node.meta.get("spec")
+            if spec is not None:
+                new_dim_order = list(int(d) for d in explicit_dim_order)
+                spec.dim_order = new_dim_order
+                spec.stride = tuple(stride_from_dim_order(spec.shape, new_dim_order))
+    # Out-variant: set the out node's spec
+    out_node = node.kwargs.get("out")
+    if not isinstance(out_node, Node):
+        return
+    spec = out_node.meta.get("spec")
+    if spec is None:
+        return
+    if _is_layout_transforming_op(node.target):
+        explicit_dim_order = node.kwargs.get("dim_order")
+        if explicit_dim_order is not None:
+            new_dim_order = list(int(d) for d in explicit_dim_order)
+            spec.dim_order = new_dim_order
+            spec.stride = tuple(stride_from_dim_order(spec.shape, new_dim_order))
+    elif node.target in _FORMAT_PRESERVING_OPS:
+        primary = _get_primary_tensor_input(node)
+        if primary is None:
+            return
+        input_val = primary.meta.get("val")
+        if not isinstance(input_val, torch.Tensor):
+            return
+        new_dim_order = dim_order_from_stride(input_val)
+        spec.dim_order = new_dim_order
+        spec.stride = tuple(stride_from_dim_order(spec.shape, new_dim_order))
+
+
 # pyre-ignore
 def make_spec(x):
     if isinstance(x, ProxyValue):
@@ -37,14 +148,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
@@ -79,18 +188,45 @@ def get_spec(x):
                         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:
+                        else:
+                            node.meta["spec"] = pytree.tree_map(make_spec, meta_val)
+                        continue
+                    # Layout-transforming ops (e.g. _to_dim_order_copy) may lack meta["val"];
+                    # ensure they get a spec from primary input + dim_order kwarg.
+                    if (
+                        "spec" not in node.meta
+                        and node.op == "call_function"
+                        and _is_layout_transforming_op(node.target)
+                    ):
+                        explicit_dim_order = node.kwargs.get("dim_order")
+                        primary = _get_primary_tensor_input(node)
+                        if explicit_dim_order is not None and primary is not None:
+                            inp_spec = primary.meta.get("spec")
+                            if isinstance(inp_spec, TensorSpec):
+                                # Use dtype from op kwarg when present (e.g. _to_dim_order_copy(..., dtype=torch.double))
+                                output_dtype = node.kwargs.get("dtype", inp_spec.dtype)
+                                node.meta["spec"] = TensorSpec(
+                                    dtype=output_dtype,
+                                    shape=inp_spec.shape,
+                                    layout=inp_spec.layout,
+                                    is_sparse=inp_spec.is_sparse,
+                                    const=inp_spec.const,
+                                    requires_grad=inp_spec.requires_grad,
+                                )
+                                node.meta["spec"].stride = tuple(
+                                    stride_from_dim_order(
+                                        inp_spec.shape, list(explicit_dim_order)
+                                    )
+                                )
+                                node.meta["spec"].dim_order = list(
+                                    int(d) for d in explicit_dim_order
+                                )
+                    if "spec" not in node.meta and meta_val is not None:
                         node.meta["spec"] = pytree.tree_map(make_spec, meta_val)
+                    _fix_out_spec_dim_order(node)
+
         return res
 
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
@@ -115,7 +251,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