diff --git a/backends/arm/operators/op_placeholder.py b/backends/arm/operators/op_placeholder.py index b5dcf3f9873..2618c9e71d3 100644 --- a/backends/arm/operators/op_placeholder.py +++ b/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 diff --git a/backends/arm/runtime/ArmBackendEthosU.cpp b/backends/arm/runtime/ArmBackendEthosU.cpp index 6d9ab6b0091..ec918433b01 100644 --- a/backends/arm/runtime/ArmBackendEthosU.cpp +++ b/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. diff --git a/backends/arm/test/misc/test_dim_order_guards.py b/backends/arm/test/misc/test_dim_order_guards.py new file mode 100644 index 00000000000..8bad1493b1c --- /dev/null +++ b/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()