Keep cat_out's nnlib input arrays off variable-length stack arrays (#24319)

backends/cadence/aot/BUCK

@@ -502,6 +502,7 @@ fbcode_target(_kind = python_unittest,
         "//executorch/backends/cadence/aot:pass_utils",
         "//executorch/backends/cadence/aot:remove_ops",
         "//executorch/backends/cadence/aot/quantizer:quantizer",
+        "//executorch/exir:lib",
         "//executorch/exir/dialects:lib",
     ],
 )

backends/cadence/aot/passes.py

@@ -21,7 +21,7 @@
 from executorch.backends.cadence.aot.remove_ops import (
     CadenceRemoveNops,
     RemoveAliasCopyOpPass,
-    RemoveCloneOpsTransformImported,
+    RemoveCloneOpsPass,
     RemoveDetachCopyPass,
     RemoveNopExpandOpPass,
     RemoveNopSliceOrViewOpPass,
@@ -114,7 +114,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
         FinalizePipeline,
         FuseSliceSameDimPass,
         RemoveAliasCopyOpPass,
-        RemoveCloneOpsTransformImported,
+        RemoveCloneOpsPass,
         RemoveDetachCopyPass,
         RemoveNopExpandOpPass,
         RemoveToOpsPass,

backends/cadence/aot/remove_ops.py

@@ -29,13 +29,12 @@
 )
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.dialects.edge._ops import EdgeOpOverload, EdgeOpOverloadPacket
 from executorch.exir.pass_base import (
     ExportedProgramPassBase,
     ExportedProgramPassResult,
     ExportPass,
     PassResult,
 )
-from executorch.exir.pass_manager import PassManager, PassType
 from executorch.exir.passes import dead_code_elimination_pass
 from torch.export import ExportedProgram
 from torch.export.graph_signature import InputKind, OutputKind
@@ -44,12 +43,32 @@
 from torch.utils import _pytree as pytree
 
 
-class RemoveCloneOpsTransformImported(ExportPass):
+class RemoveCloneOpsPass(RemoveCloneOpsTransform):
+    """
+    Remove clones, including those that only request a memory format.
+
+    Cadence lowers with ``_skip_dim_order=True``: every edge tensor is
+    contiguous and a memory format has no runtime representation, so such a
+    clone is a plain copy. Left in the graph, the ExecuTorch emitter rejects
+    any format other than contiguous or preserve.
+
+    Before:
+        clone(x, memory_format=torch.channels_last)
+
+    After:
+        x
+    """
+
+    def __init__(self) -> None:
+        super().__init__(eliminate_quant_dequant_pairs=False)
+
+    def _is_non_identity_clone(self, node: torch.fx.Node) -> bool:
+        if node.target == exir_ops.edge.aten.clone.default:
+            return False
+        return super()._is_non_identity_clone(node)
+
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
-        finalize_passes: List[PassType] = [
-            RemoveCloneOpsTransform(eliminate_quant_dequant_pairs=False),
-        ]
-        result = PassManager(passes=finalize_passes)(graph_module)
+        result = super().call(graph_module)
         dead_code_elimination_pass(result.graph_module)
         return result
 
@@ -911,7 +930,7 @@
         FuseTransposeOrPermuteOpPairsPass,
         RemoveSqueezeViewBeforeElementwiseOps,
         RemoveCatFromSliceCopyPass,
-        RemoveCloneOpsTransformImported,
+        RemoveCloneOpsPass,
     ]
 
 

backends/cadence/aot/tests/test_remove_ops_passes.py

@@ -20,7 +20,7 @@
     RemoveAliasCopyOpPass,
     RemoveBranchedQuantDequant,
     RemoveCatFromSliceCopyPass,
-    RemoveCloneOpsTransformImported,
+    RemoveCloneOpsPass,
     RemoveContiguousOpPass,
     RemoveDetachCopyPass,
     RemoveNopAddOpPass,
@@ -38,6 +38,7 @@
 )
 from executorch.backends.cadence.aot.typing_stubs import expand
 from executorch.backends.test.graph_builder import GraphBuilder
+from executorch.exir import EdgeCompileConfig, to_edge
 from executorch.exir.dialects._ops import ops as exir_ops
 from pyre_extensions import none_throws
 
@@ -282,12 +283,66 @@ def test_remove_clone(self) -> None:
         clone = builder.call_operator(op=exir_ops.edge.aten.clone.default, args=(x,))
         builder.output([clone])
         original = builder.get_graph_module()
-        p = RemoveCloneOpsTransformImported()
+        p = RemoveCloneOpsPass()
         graph_after_passes = cast(PassResult, p(original)).graph_module
         self.assertEqual(
             count_node(graph_after_passes, torch.ops.aten.clone.default), 0
         )
 
+    @expand(
+        [
+            ["channels_last", torch.channels_last],
+            ["contiguous", torch.contiguous_format],
+            ["preserve", torch.preserve_format],
+        ]
+    )
+    def test_remove_clone_with_memory_format(
+        self, _name: str, memory_format: torch.memory_format
+    ) -> None:
+        builder = GraphBuilder()
+        x = builder.placeholder("x", torch.randn([1, 8, 4, 4], dtype=torch.float32))
+        clone = builder.call_operator(
+            op=exir_ops.edge.aten.clone.default,
+            args=(x,),
+            kwargs={"memory_format": memory_format},
+        )
+        builder.output([clone])
+        original = builder.get_graph_module()
+        p = RemoveCloneOpsPass()
+        result = cast(PassResult, p(copy.deepcopy(original)))
+        self.assertTrue(result.modified)
+        self.assertEqual(
+            count_node(result.graph_module, exir_ops.edge.aten.clone.default), 0
+        )
+        validate(
+            original,
+            result.graph_module,
+            (torch.randn(1, 8, 4, 4),),
+            "RemoveCloneOpsPass",
+        )
+
+    def test_remove_clone_lets_channels_last_program_emit(self) -> None:
+        # rlx_arbitrary_image_stylization: a CPU-side
+        # clone(memory_format=torch.channels_last) failed in the emitter.
+        class ChannelsLastCopy(torch.nn.Module):
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                return x.contiguous(memory_format=torch.channels_last) + 1.0
+
+        ep = torch.export.export(ChannelsLastCopy(), (torch.randn(1, 8, 4, 4),))
+        config = EdgeCompileConfig(_skip_dim_order=True)
+        with self.assertRaisesRegex(Exception, "memory_format"):
+            to_edge(copy.deepcopy(ep), compile_config=config).to_executorch()
+
+        edge = to_edge(ep, compile_config=config).transform([RemoveCloneOpsPass()])
+        self.assertEqual(
+            count_node(
+                edge.exported_program().graph_module,
+                exir_ops.edge.aten.clone.default,
+            ),
+            0,
+        )
+        self.assertGreater(len(edge.to_executorch().buffer), 0)
+
     def test_remove_contiguous(self) -> None:
         builder = GraphBuilder()
         x = builder.placeholder("x", torch.randn([3, 5], dtype=torch.float32))

backends/cadence/fusion_g3/operators/op_cat.cpp

@@ -1,3 +1,3 @@
 /*
  * Copyright (c) Meta Platforms, Inc. and affiliates.
  * All rights reserved.
@@ -33,6 +33,12 @@
 namespace G3 {
 namespace native {
 
+namespace {
+// Upper bound on inputs for the xa_nn_cat path, whose per-input arrays live on
+// the DSP thread stack. Larger cats take the generic loop.
+constexpr size_t kMaxNnlibCatInputs = 16;
+} // namespace
+
 Tensor& cat_out(
     KernelRuntimeContext& ctx,
     ArrayRef<Tensor> tensors,
@@ -42,7 +48,7 @@
     dim += out.dim();
   }
 
-  int kTensorDimensionLimit = executorch::runtime::kTensorDimensionLimit;
+  constexpr int kTensorDimensionLimit = executorch::runtime::kTensorDimensionLimit;
 
 #ifdef OP_ARG_CHECK
 
@@ -76,34 +82,6 @@
     return out;
   }
 
-  const signed char* inp_tensors[tensors.size()];
-  const int* inp_tensors_shapes[tensors.size()];
-
-  int inp_shapes_size[tensors.size()];
-
-  int temp_sizes[tensors.size()][kTensorDimensionLimit];
-  ArrayRef<Tensor::SizesType> temp_size;
-
-  for (int i = 0; i < tensors.size(); i++) {
-    inp_tensors[i] = tensors[i].const_data_ptr<signed char>();
-    temp_size = tensors[i].sizes();
-
-    for (int j = 0; j < temp_size.size(); j++) {
-      temp_sizes[i][j] = temp_size[j];
-    }
-    inp_tensors_shapes[i] = temp_sizes[i]; // input shapes
-    inp_shapes_size[i] = temp_size.size(); // number of input dimensions
-  }
-
-  signed char* out_data = out.mutable_data_ptr<signed char>();
-
-  const ArrayRef<Tensor::SizesType> out_size = out.sizes();
-  int out_shapes[kTensorDimensionLimit];
-  for (int i = 0; i < out_size.size(); i++) // output shapes
-  {
-    out_shapes[i] = out_size[i];
-  }
-
   bool all_same_dtype = true;
   for (int i = 0; i < tensors.size(); i++) {
     if (out.scalar_type() != tensors[i].scalar_type()) {
@@ -120,7 +98,33 @@
       out.scalar_type() == ScalarType::Byte ||
       out.scalar_type() == ScalarType::Float;
 
-  if (all_same_dtype && supported_dtype) {
+  if (all_same_dtype && supported_dtype &&
+      tensors.size() <= kMaxNnlibCatInputs) {
+    const signed char* inp_tensors[kMaxNnlibCatInputs];
+    const int* inp_tensors_shapes[kMaxNnlibCatInputs];
+    int inp_shapes_size[kMaxNnlibCatInputs];
+    int temp_sizes[kMaxNnlibCatInputs][kTensorDimensionLimit];
+
+    for (int i = 0; i < tensors.size(); i++) {
+      inp_tensors[i] = tensors[i].const_data_ptr<signed char>();
+      const ArrayRef<Tensor::SizesType> temp_size = tensors[i].sizes();
+
+      for (int j = 0; j < temp_size.size(); j++) {
+        temp_sizes[i][j] = temp_size[j];
+      }
+      inp_tensors_shapes[i] = temp_sizes[i]; // input shapes
+      inp_shapes_size[i] = temp_size.size(); // number of input dimensions
+    }
+
+    signed char* out_data = out.mutable_data_ptr<signed char>();
+
+    const ArrayRef<Tensor::SizesType> out_size = out.sizes();
+    int out_shapes[kTensorDimensionLimit];
+    for (int i = 0; i < out_size.size(); i++) // output shapes
+    {
+      out_shapes[i] = out_size[i];
+    }
+
     XT_KERNEL_CHECK(
         ctx,
         out,