[EXIR] Register _clone_dim_order op and map aten.clone (#13735)

### Summary
This is PR 2 of 3 implementing a dim order aware clone op. 

This PR registers the new `_clone_dim_order` op and maps `aten.clone` to
`dim_order_ops._clone_dim_order` in EXIR during export to preserve
memory layout changes (contiguous/channels_last). It also updates the
Core ML, ARM, and Qualcomm backends to handle the new clone op.

Related PRs:
- PR 1: [#12974](#12974) - Add
`_clone_dim_order` portable kernel
- PR 3: [#12976](#12976) -
Update RemoveCloneOpsTransform to be dim order aware

Fixes #12645  

### Test plan
- Operator level tests to verify kernel behavior for layout preservation
and changes.
- Graph level checks to confirm that clone mapping occurs.
- End to end tests to validate that functional clone behavior is
unchanged.
- Backend tests to ensure clone semantics are preserved.

All tests pass via:
`python -m unittest exir.tests.test_memory_format_ops_pass`
`python -m unittest backends.apple.coreml.test.test_torch_ops`
`pytest backends/arm/test/ops/test_clone.py`
`pytest backends/arm/test/passes/test_remove_clone_pass.py`

---------

Co-authored-by: Gasoonjia <[email protected]>
Co-authored-by: Digant Desai <[email protected]>

Author: 3 people

backends/apple/coreml/compiler/torch_ops.py

@@ -15,6 +15,7 @@
 from coremltools.converters.mil.frontend.torch.ops import (
     _get_inputs,
     _get_kwinputs,
+    noop,
     NUM_TO_NUMPY_DTYPE,
     NUM_TO_TORCH_DTYPE,
     split,
@@ -91,6 +92,28 @@ def _to_dim_order_copy(context, node):
         to(context, node)
 
 
+@register_torch_op(
+    torch_alias=[
+        "dim_order_ops::_clone_dim_order",
+        "dim_order_ops._clone_dim_order",
+    ],
+    override=False,
+)
+def _clone_dim_order(context, node):
+    dim_order = _get_kwinputs(context, node, "dim_order", default=[None])[0]
+    node.kwinputs.pop("dim_order")
+
+    # In CoreML, dim_order.val will be a ndarray, so we convert it to a list to check memory format.
+    dim_order = [int(d) for d in dim_order.val]
+    memory_format = get_memory_format(dim_order)
+    assert (
+        memory_format == _torch.contiguous_format
+    ), "Only contiguous memory format is supported in CoreML"
+
+    # Since CoreML only supports contiguous format, no dim_order preservation is needed. Treat this as a no-op clone.
+    noop(context, node)
+
+
 # https://github.com/apple/coremltools/pull/2558
 @register_torch_op(
     torch_alias=["torchao::dequantize_affine", "torchao.dequantize_affine"],

backends/apple/coreml/test/test_torch_ops.py

@@ -268,6 +268,28 @@ def test_dequantize_codebook_embedding_per_grouped_row(self):
         et_prog = delegated_program.to_executorch()
         self._compare_outputs(et_prog, model, example_inputs)
 
+    def test__clone_dim_order_contiguous(self):
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                return torch.ops.dim_order_ops._clone_dim_order(
+                    x, dim_order=[0, 1, 2, 3]
+                )
+
+        model, example_inputs = Model(), (torch.randn(1, 3, 8, 8),)
+        ep = torch.export.export(model, example_inputs)
+        delegated_program = executorch.exir.to_edge_transform_and_lower(
+            ep,
+            partitioner=[self._coreml_partitioner()],
+        )
+        for node in delegated_program.exported_program().graph.nodes:
+            if node.op == "call_function":
+                assert node.target.__name__ in [
+                    "executorch_call_delegate",
+                    "getitem",
+                ], f"Got unexpected node target after delegation: {node.target.__name__}"
+        et_prog = delegated_program.to_executorch()
+        self._compare_outputs(et_prog, model, example_inputs)
+
 
 if __name__ == "__main__":
     test_runner = TestTorchOps()
@@ -280,3 +302,4 @@ def test_dequantize_codebook_embedding_per_grouped_row(self):
     test_runner.test_dequantize_codebook_linear_per_grouped_row()
     test_runner.test_dequantize_codebook_embedding_per_grouped_col()
     test_runner.test_dequantize_codebook_embedding_per_grouped_row()
+    test_runner.test__clone_dim_order_contiguous()

backends/arm/_passes/remove_clone_pass.py

@@ -18,7 +18,7 @@ class RemoveClonePass(ExportPass):
     """Remove all clones from graph_module"""
 
     def call_operator(self, op, args, kwargs, meta):
-        if op != exir_ops.edge.aten.clone.default:
+        if op != exir_ops.edge.dim_order_ops._clone_dim_order.default:
             return super().call_operator(op, args, kwargs, meta)
 
         if len(args) != 1:

backends/arm/operator_support/clone_support.py

@@ -5,6 +5,7 @@
 
 import logging
 
+import torch
 import torch.fx as fx
 from executorch.backends.arm.operator_support.tosa_supported_operators import (
     register_tosa_support_check,
@@ -18,7 +19,7 @@
 
 @register_tosa_support_check
 class CloneSupported(SupportedTOSAOperatorCheck):
-    targets = [exir_ops.edge.aten.clone.default]
+    targets = [exir_ops.edge.dim_order_ops._clone_dim_order.default]
 
     tosa_specs = [
         TosaSpecification.create_from_string("TOSA-1.0+INT"),
@@ -28,10 +29,62 @@ class CloneSupported(SupportedTOSAOperatorCheck):
     def is_node_tosa_supported(
         self, node: fx.Node, tosa_spec: TosaSpecification
     ) -> bool:
+        if node.target not in self.targets:
+            self.reporter.report_reject(node, f"Target {node.target} is not supported.")
+            return False
 
         input_node = node.args[0]
         if not isinstance(input_node, fx.Node):
             self.reporter.report_reject(node, "Non tensor clones are not supported")
             return False
 
+        # Check input node
+        if len(node.all_input_nodes) != 1:
+            self.reporter.report_reject(
+                node, f"Expected 1 input node, got {len(node.all_input_nodes)}"
+            )
+            return False
+
+        input_val = node.all_input_nodes[0].meta["val"]
+        if not isinstance(input_val, torch._subclasses.FakeTensor):
+            self.reporter.report_reject(node, "Expected input to be a FakeTensor.")
+            return False
+
+        input_dtype = input_val.dtype
+
+        # Check output node
+        output_val = node.meta["val"]
+        if not isinstance(output_val, torch._subclasses.FakeTensor):
+            self.reporter.report_reject(node, "Expected output to be a FakeTensor.")
+            return False
+
+        if output_val.dtype != input_dtype:
+            self.reporter.report_reject(
+                node,
+                f"Input dtype {input_val.dtype} does not match {output_val.dtype}.",
+            )
+            return False
+
+        # Check memory format
+        if "memory_format" in node.kwargs:
+            if node.kwargs["memory_format"] in (torch.preserve_format,):
+                self.reporter.report_reject(
+                    node,
+                    f"Argument 'memory_format' is not supported for "
+                    f"{node.target} right now.",
+                )
+                return False
+
+        # Check dim_order
+        if "dim_order" in node.kwargs:
+            dim_order = node.kwargs["dim_order"]
+            # pyre-ignore[6]
+            if dim_order != list(range(len(dim_order))):  # type: ignore[arg-type]
+                self.reporter.report_reject(
+                    node,
+                    f"Argument {dim_order=} is not supported for "
+                    f"{node.target} right now.",
+                )
+                return False
+
         return True

backends/arm/test/misc/test_partition_decomposed_quantized_ops.py

@@ -38,7 +38,7 @@
 ]
 linear_residual_exir_op: list[str] = [
     "executorch_exir_dialects_edge__ops_aten_gelu_default",
-    "executorch_exir_dialects_edge__ops_aten_clone_default",
+    "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default",
     "executorch_exir_dialects_edge__ops_aten_linear_default",
     "executorch_exir_dialects_edge__ops_aten_add_Tensor",
 ]

backends/arm/test/ops/test_clone.py

@@ -19,7 +19,7 @@
 )
 
 aten_op = "torch.ops.aten.clone.default"
-exir_op = "executorch_exir_dialects_edge__ops_aten_clone_default"
+exir_op = "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default"
 
 input_t = Tuple[torch.Tensor]
 

backends/arm/test/passes/test_remove_clone_pass.py

@@ -35,9 +35,11 @@ def test_remove_clone_tosa_INT():
         module.get_inputs(),
         quantize=True,
         ops_before_pass={
-            "executorch_exir_dialects_edge__ops_aten_clone_default": 1,
+            "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 1,
         },
-        ops_not_after_pass=["executorch_exir_dialects_edge__ops_aten_clone_default"],
+        ops_not_after_pass=[
+            "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default"
+        ],
         pass_list=[RemoveClonePass],
     )
     pipeline.run()

backends/qualcomm/_passes/convert_bmm_to_matmul.py

@@ -21,7 +21,7 @@ class ConvertBmmToMatmul(ExportPass):
 
     view_copy = exir_ops.edge.aten.view_copy.default
     expand_copy = exir_ops.edge.aten.expand_copy.default
-    clone = exir_ops.edge.aten.clone.default
+    clone = exir_ops.edge.dim_order_ops._clone_dim_order.default
     bmm = exir_ops.edge.aten.bmm.default
     matmul = exir_ops.edge.aten.matmul.default
     patterns = [

backends/qualcomm/_passes/remove_redundancy.py

@@ -19,7 +19,7 @@ def __init__(self, quantization_capture=False):
         self.redundant_ops_general = {
             torch.clone: self._default_condition,
             torch.ops.aten.clone.default: self._default_condition,
-            exir_ops.edge.aten.clone.default: self._default_condition,
+            exir_ops.edge.dim_order_ops._clone_dim_order.default: self._default_condition,
             torch.ops.aten.alias.default: self._default_condition,
             exir_ops.edge.aten.alias.default: self._default_condition,
             exir_ops.edge.aten.alias_copy.default: self._default_condition,

backends/qualcomm/partition/common_defs.py

@@ -10,7 +10,7 @@
 from executorch.exir.dialects._ops import ops as exir_ops
 
 not_supported_operator = [
-    exir_ops.edge.aten.clone.default,
+    exir_ops.edge.dim_order_ops._clone_dim_order.default,
     exir_ops.edge.quantized_decomposed.embedding_4bit.dtype,
 ]