Update on "Arm backend: Add 16A8W support and test for mul operation"

Add 16A8W quantization support and test for the mul operation in ExecutorTorch ARM backend.

This follows the pattern established for linear operations, extending int16 support to mul operations.

Changes:
- Add INT16 dtype validation support in op_mul.py
- Add test_mul_tensor_16a8w_tosa_INT test function
- Enable test_mul.py in test targets configuration

The 16A8W configuration uses 16-bit activations with 8-bit weights, enabling higher precision for activations while maintaining weight efficiency.

Differential Revision: [D80510628](https://our.internmc.facebook.com/intern/diff/D80510628/)

cc digantdesai freddan80 per zingo oscarandersson8218

[ghstack-poisoned]

Author: Ninja91

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/arm_pass_manager.py

@@ -91,7 +91,8 @@
     UnsqueezeBeforeRepeatPass,
     UnsqueezeScalarPlaceholdersPass,
 )
-from executorch.backends.arm.tosa_specification import (
+
+from executorch.backends.arm.tosa.specification import (
     TosaLoweringContext,
     TosaSpecification,
 )

backends/arm/_passes/decompose_meandim_pass.py

@@ -9,7 +9,6 @@
 import torch
 from executorch.backends.arm._passes import ArmPass
 from executorch.backends.arm._passes.arm_pass_utils import get_node_arg
-from executorch.backends.arm.operator_support.pool_2d_support import AvgPool2dSupported
 from executorch.exir.backend.utils import WhyNoPartitionReporter
 from executorch.exir.dialects._ops import ops as exir_ops
 
@@ -67,6 +66,11 @@ def __init__(self, graph_module, tosa_spec):
         super().__init__()
         self._graph_module = graph_module
         self._tosa_spec = tosa_spec
+        # Lazy import to avoid circular dependency with operator_support
+        from executorch.backends.arm.operator_support.pool_2d_support import (
+            AvgPool2dSupported,
+        )
+
         self._avg_pool_checker = AvgPool2dSupported(
             self._tosa_spec, WhyNoPartitionReporter()
         )

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/arm_backend.py

@@ -13,9 +13,7 @@
 from enum import Enum
 from typing import List, Optional
 
-from executorch.backends.arm.tosa_specification import (  # type: ignore[import-not-found]
-    TosaSpecification,
-)
+from executorch.backends.arm.tosa import TosaSpecification
 
 from executorch.exir.backend.compile_spec_schema import (  # type: ignore[import-not-found]
     CompileSpec,

backends/arm/ethosu/backend.py

@@ -16,7 +16,7 @@
 
 from executorch.backends.arm.arm_vela import vela_compile
 
-from executorch.backends.arm.tosa_backend import TOSABackend
+from executorch.backends.arm.tosa.backend import TOSABackend
 from executorch.exir.backend.backend_details import BackendDetails, PreprocessResult
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 from torch.export.exported_program import ExportedProgram

backends/arm/ethosu/partitioner.py

@@ -11,7 +11,7 @@
     is_ethosu,
 )  # usort: skip
 from executorch.backends.arm.ethosu import EthosUBackend
-from executorch.backends.arm.tosa_partitioner import TOSAPartitioner
+from executorch.backends.arm.tosa.partitioner import TOSAPartitioner
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 from executorch.exir.backend.partitioner import DelegationSpec
 from torch.fx.passes.operator_support import OperatorSupportBase

backends/arm/operator_support/clone_support.py

@@ -5,20 +5,21 @@
 
 import logging
 
+import torch
 import torch.fx as fx
 from executorch.backends.arm.operator_support.tosa_supported_operators import (
     register_tosa_support_check,
     SupportedTOSAOperatorCheck,
 )
-from executorch.backends.arm.tosa_specification import TosaSpecification
+from executorch.backends.arm.tosa import TosaSpecification
 from executorch.exir.dialects._ops import ops as exir_ops
 
 logger = logging.getLogger(__name__)
 
 
 @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/operator_support/convolution_support.py

@@ -11,7 +11,7 @@
     register_tosa_support_check,
     SupportedTOSAOperatorCheck,
 )
-from executorch.backends.arm.tosa_specification import TosaSpecification
+from executorch.backends.arm.tosa import TosaSpecification
 
 from executorch.exir.dialects._ops import ops as exir_ops