Fix channels-last portable fallback by inserting explicit dim-order legalization (#20878)

Fixes #16507 #11086 

I kept the pass and regressions in separate files to keep the change
cleaner. A more optimized followup is probably worth doing if more
issues like these appear.

Added separate test file
`exir/tests/test_legalize_portable_dim_order_pass.py`

cc @JacobSzwejbka @angelayi @larryliu0820 @lucylq

Author: arkadip-maitra

exir/passes/__init__.py

@@ -40,6 +40,9 @@
 from executorch.exir.passes.insert_write_back_for_buffers_pass import (
     insert_write_back_for_buffers_pass,
 )
+from executorch.exir.passes.legalize_portable_dim_order_pass import (
+    LegalizePortableDimOrderPass,
+)
 from executorch.exir.passes.memory_format_ops_pass import MemoryFormatOpsPass
 from executorch.exir.passes.memory_planning_pass import MemoryPlanningPass
 from executorch.exir.passes.normalize_transpose_pass import NormalizeTransposePass
@@ -76,6 +79,7 @@
     "OpReplacePass",
     "ToDevicePass",
     "EdgeToBackendOpsPass",
+    "LegalizePortableDimOrderPass",
     "MemoryFormatOpsPass",
     "MemoryPlanningPass",
     "HintBasedSymShapeEvalPass",

exir/passes/legalize_portable_dim_order_pass.py

@@ -0,0 +1,69 @@
+# 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.
+
+from typing import Optional
+
+import executorch.exir.passes.dim_order_ops_registry  # noqa: F401
+
+import torch
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.dim_order_utils import get_dim_order
+from executorch.exir.pass_base import ExportPass, NodeMetadata, ProxyValue
+
+
+class LegalizePortableDimOrderPass(ExportPass):
+    """Insert contiguous dim-order copies before portable default-layout ops.
+
+    Runs during ``edge -> executorch`` lowering on leftover edge ops whose
+    portable kernels require default dim order on their primary input.
+    """
+
+    _copy_op = exir_ops.edge.dim_order_ops._to_dim_order_copy.default
+    _target_ops = {
+        exir_ops.edge.aten._adaptive_avg_pool2d.default,
+        exir_ops.edge.aten.avg_pool2d.default,
+        exir_ops.edge.aten.expand_copy.default,
+        exir_ops.edge.aten.mean.dim,
+        exir_ops.edge.aten.sum.dim_IntList,
+    }
+
+    def call_operator(self, op, args, kwargs, meta):
+        if op not in self._target_ops:
+            return super().call_operator(op, args, kwargs, meta)
+
+        input_arg = args[0]
+        if isinstance(input_arg, ProxyValue) and input_arg.is_tensor():
+            input_tensor: Optional[torch.Tensor] = input_arg.to_tensor()
+            input_meta = NodeMetadata(input_arg.node.meta)
+        elif isinstance(input_arg, torch.Tensor):
+            input_tensor = input_arg
+            input_meta = meta.copy()
+        else:
+            input_tensor = None
+            input_meta = meta.copy()
+
+        if input_tensor is None or tuple(
+            int(d) for d in input_tensor.dim_order()
+        ) == tuple(range(input_tensor.dim())):
+            return super().call_operator(op, args, kwargs, meta)
+
+        contiguous_dim_order = get_dim_order(
+            torch.contiguous_format, input_tensor.dim()
+        )
+        assert contiguous_dim_order is not None
+
+        legalized_input = super().call_operator(
+            self._copy_op,
+            (input_arg,),
+            {"dim_order": contiguous_dim_order},
+            input_meta,
+        )
+        return super().call_operator(
+            op,
+            (legalized_input, *args[1:]),
+            kwargs,
+            meta,
+        )

exir/program/_program.py

@@ -45,6 +45,7 @@
     convert_constant_dim_order_pass,
     dead_code_elimination_pass,
     EdgeToBackendOpsPass,
+    LegalizePortableDimOrderPass,
     MemoryFormatOpsPass,
     OpReplacePass,
     remove_unused_parameters_pass,
@@ -405,12 +406,14 @@ def __init__(
         self,
         exported_program: ExportedProgram,
         after_to_edge_passes: bool,
+        compile_config: Optional[EdgeCompileConfig] = None,
     ):
         self.exported_program = exported_program
 
         # Add a flag to denote whehter to_edge is called on this program
         # to detect misusage of directly calling to_executorch without to_edge
         self.after_to_edge_passes = after_to_edge_passes
+        self.compile_config = compile_config or EdgeCompileConfig()
 
     def transform(self, *passes: PassType) -> "ExirExportedProgram":
         self.exported_program = _transform(self.exported_program, *passes)
@@ -441,7 +444,9 @@ def to_executorch(
             raise RuntimeError("Must run to_edge before to_executorch.")
         config = config or ExecutorchBackendConfig()
         new_gm = self.exported_program.graph_module
-        for p in edge_to_executorch_passes(config):
+        for p in edge_to_executorch_passes(
+            config, edge_compile_config=self.compile_config
+        ):
             new_gm_res = p(new_gm)
             assert new_gm_res is not None
             new_gm = new_gm_res.graph_module
@@ -478,6 +483,7 @@ def __deepcopy__(
         new_eep = ExirExportedProgram(
             copy.deepcopy(self.exported_program, memo),
             self.after_to_edge_passes,
+            copy.deepcopy(self.compile_config, memo),
         )
         return new_eep
 
@@ -673,6 +679,7 @@ def _to_edge(ep, config: EdgeCompileConfig) -> "ExirExportedProgram":
                 ],
             ),
             False,
+            config,
         )
     pre_op_replace_passes, post_op_replace_passes = _get_aten_to_edge_passes(config)
 
@@ -753,7 +760,9 @@ def pre_memory_planning_passes(
 
 
 def edge_to_executorch_passes(
-    config: ExecutorchBackendConfig, name: Optional[str] = None
+    config: ExecutorchBackendConfig,
+    name: Optional[str] = None,
+    edge_compile_config: Optional[EdgeCompileConfig] = None,
 ) -> List[PassType]:
     """
     Returns a list of passes to lower from edge to executorch.
@@ -766,6 +775,7 @@ def edge_to_executorch_passes(
     else:
         device_cfg = propagate_device_config
 
+    edge_compile_config = edge_compile_config or EdgeCompileConfig()
     passes: List[PassType] = [
         # ExecuTorch backend ops are unable to handle unbacked symints. So after
         # this pass, passes cannot be Interpreter-based, because it will fail if
@@ -781,6 +791,9 @@ def edge_to_executorch_passes(
         RemoveGraphAssertsPass(),
     ] + pre_memory_planning_passes(config, name)
 
+    if not edge_compile_config._skip_dim_order:
+        passes.insert(len(config.passes), LegalizePortableDimOrderPass())
+
     return passes
 
 
@@ -1702,7 +1715,9 @@ def to_executorch(  # noqa (FLAKE8) C901
             program = unsafe_remove_auto_functionalized_pass(program)
             gm, new_signature = insert_write_back_for_buffers_pass(program)
             new_gm = program.graph_module
-            for p in edge_to_executorch_passes(config, name):
+            for p in edge_to_executorch_passes(
+                config, name, edge_compile_config=self.compile_config
+            ):
                 new_gm_res = p(new_gm)
                 assert new_gm_res is not None
                 new_gm = new_gm_res.graph_module