Support dim order in Arm backend (#5576)

Summary:
Add both ahead-of-time and runtime guards to make sure we don't accept inputs with any other memory-format than contiguous format.

Change-Id: I9e29badbcf238d458e12f0d62394abae66d421b7

Pull Request resolved: #5576

Reviewed By: mergennachin

Differential Revision: D63638653

Pulled By: digantdesai

fbshipit-source-id: 9eea86836919eacd5578ef024cd43d039fd40123

Author: oscarandersson8218

backends/arm/operators/op_placeholder.py

@@ -28,6 +28,13 @@ def process_inputs(
     tosa_graph: ts.TosaSerializer,
 ):
     """Serialize an input node"""
+    # inputs need to be in default dim_order (contiguous memory format)
+    meta = node.meta["val"]
+    if meta.dim_order() != tuple(range(meta.dim())):
+        raise RuntimeError(
+            f"Arm backend only supports contiguous memory format for inputs. "
+            f"Expected dim_order: {tuple(range(meta.dim()))}, but got: {meta.dim_order()} for node {node.name}"
+        )
     inputs = [TosaArg(node)]
     input_shape = inputs[0].shape
     input_dim_order = inputs[0].dim_order

backends/arm/runtime/ArmBackendEthosU.cpp

@@ -19,6 +19,7 @@
 #include "executorch/runtime/backend/interface.h"
 #include "executorch/runtime/core/error.h"
 #include "executorch/runtime/core/evalue.h"
+#include "executorch/runtime/core/exec_aten/util/dim_order_util.h"
 #include "executorch/runtime/core/exec_aten/util/scalar_type_util.h"
 
 using namespace std;
@@ -144,6 +145,15 @@ class ArmBackend final : public ::executorch::runtime::BackendInterface {
             toString(tensor_in.scalar_type()));
         return Error::InvalidProgram;
       }
+      supported = is_contiguous_dim_order(
+          tensor_in.dim_order().data(), tensor_in.dim());
+      if (!supported) {
+        ET_LOG(
+            Error,
+            "Input %d expected contiguous dim_order, but got non-contiguous dim_order",
+            i);
+        return Error::InvalidProgram;
+      }
 
       // Select a compatible copy routine including checking for input layouts
       // which require permutation.

backends/arm/test/misc/test_dim_order_guards.py

@@ -0,0 +1,58 @@
+# Copyright 2024 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+
+import pytest
+
+import torch
+from executorch.backends.arm.test import common
+
+from executorch.backends.arm.test.tester.arm_tester import ArmTester
+
+
+class Conv2D(torch.nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.conv2d = torch.nn.Conv2d(in_channels=2, out_channels=3, kernel_size=(3, 3))
+
+    def forward(self, x):
+        return self.conv2d(x.to(memory_format=torch.channels_last))
+
+    def get_inputs(self):
+        return (torch.randn(1, 2, 20, 20),)
+
+
+class TestDimOrderGuards(unittest.TestCase):
+
+    def test_tosa_MI_pipeline(self):
+        module = Conv2D()
+        tester = (
+            ArmTester(
+                module,
+                example_inputs=module.get_inputs(),
+                compile_spec=common.get_tosa_compile_spec(),
+            )
+            .export()
+            .to_edge()
+        )
+        with pytest.raises(RuntimeError):
+            tester.partition()
+
+    def test_tosa_BI_pipeline(self):
+        module = Conv2D()
+        tester = (
+            ArmTester(
+                module,
+                example_inputs=module.get_inputs(),
+                compile_spec=common.get_tosa_compile_spec(),
+            )
+            .quantize()
+            .export()
+            .to_edge()
+        )
+        with pytest.raises(RuntimeError):
+            tester.partition()