diff --git a/backends/native/serialization/graph_serialize.py b/backends/native/serialization/graph_serialize.py index 0d89842cabd..1f9a333c871 100644 --- a/backends/native/serialization/graph_serialize.py +++ b/backends/native/serialization/graph_serialize.py @@ -203,7 +203,7 @@ def _dim_order(t: torch.Tensor) -> list[int]: strides = tuple(t.stride()) sizes = list(t.shape) # dim_order_from_stride handles symbolic strides and rejects stride-0 layouts. - dim_order = [int(d) for d in dim_order_from_stride(strides)] + dim_order = [int(d) for d in dim_order_from_stride(strides, tuple(sizes))] expected = stride_from_dim_order(sizes, dim_order) for i in range(ndim): # A size-1 dim only ever indexes 0, so its stride is arbitrary and need not diff --git a/backends/native/test/test_serialize.py b/backends/native/test/test_serialize.py index edb1c538c89..0d833968567 100644 --- a/backends/native/test/test_serialize.py +++ b/backends/native/test/test_serialize.py @@ -554,10 +554,9 @@ def test_permuted_contiguous(self): self.assertEqual(_dim_order(t), [0, 2, 1]) def test_channels_last_with_size_one_channel(self): - # channels-last leaves the size-1 channel with an arbitrary stride, which - # must not be treated as a non-expressible layout. + # A size-1 channel must not make channels-last non-expressible. t = torch.randn(2, 1, 3, 4).to(memory_format=torch.channels_last) - self.assertEqual(_dim_order(t), [0, 2, 1, 3]) + self.assertEqual(_dim_order(t), [0, 2, 3, 1]) def test_sliced_layout_raises(self): t = torch.randn(4, 8)[:, :4] diff --git a/backends/nxp/tests/models/test_mlperf_tiny_keyword_spotting.py b/backends/nxp/tests/models/test_mlperf_tiny_keyword_spotting.py index ab56849e970..22fb84d66be 100644 --- a/backends/nxp/tests/models/test_mlperf_tiny_keyword_spotting.py +++ b/backends/nxp/tests/models/test_mlperf_tiny_keyword_spotting.py @@ -29,11 +29,11 @@ BOUNDS_MSE = { "PTQ": { - "channels-last": np.inf, + "channels-last": 3.5e-4, "channels-first": 5.5e-7, }, "QAT": { - "channels-last": np.inf, + "channels-last": 3.5e-4, "channels-first": 3.3e-5, }, } @@ -45,19 +45,7 @@ def reseed_model_per_test_run(): np.random.seed(23) -@pytest.mark.parametrize( - "channels_last", - [ - False, - pytest.param( - True, - marks=pytest.mark.xfail( - reason="EIEX-1082, don't forget to readjust bounds when it start working", - strict=True, - ), - ), - ], -) +@pytest.mark.parametrize("channels_last", [False, True]) def test_mlperf_tiny_kws_mse_cpu_vs_npu(mocker, request, channels_last, use_qat): # approx. 5 samples per class num_samples = 60 @@ -87,16 +75,6 @@ def test_mlperf_tiny_kws_mse_cpu_vs_npu(mocker, request, channels_last, use_qat) partial(kws.train_model_fn, channels_last=channels_last) if use_qat else None ) - # This model does not work in channels-last format when running with portable kernels. - # See more information below. - # Github issue: https://github.com/pytorch/executorch/issues/22520 - # NXP internal issue ID: EIEX-1074 - ref_model = ( - ReferenceModel.QUANTIZED_EDGE_PYTHON - if channels_last - else ReferenceModel.QUANTIZED_EXECUTORCH_CPP - ) - lower_run_compare( model, [input_spec], @@ -105,7 +83,7 @@ def test_mlperf_tiny_kws_mse_cpu_vs_npu(mocker, request, channels_last, use_qat) dataset_creator=dataset_creator, output_comparator=comparator, mocker=mocker, - reference_model=ref_model, + reference_model=ReferenceModel.QUANTIZED_EXECUTORCH_CPP, use_qat=use_qat, train_fn=train_fn, ) diff --git a/exir/emit/_emitter.py b/exir/emit/_emitter.py index 4e3baf26592..e78b860ac3b 100644 --- a/exir/emit/_emitter.py +++ b/exir/emit/_emitter.py @@ -2034,7 +2034,7 @@ def _is_buffer(node: Node, graph_signature: ExportGraphSignature) -> bool: spec.storage = real_tensor.untyped_storage() spec.stride = real_tensor.stride() - spec.dim_order = dim_order_from_stride(spec.stride) + spec.dim_order = dim_order_from_stride(spec.stride, tuple(spec.shape)) # User inputs and mutable buffers are not constants, other buffers or parameters are. if initialize_buffer and is_mutable_buffer: spec.const = True diff --git a/exir/tensor.py b/exir/tensor.py index 199c5adfafe..50fcd46fa53 100644 --- a/exir/tensor.py +++ b/exir/tensor.py @@ -49,7 +49,9 @@ def contiguous_stride_from_shape(shape: torch.Size) -> Tuple[int]: return tuple(reversed(strides)) -def dim_order_from_stride(stride: Tuple[int]) -> Tuple[bytes]: +def dim_order_from_stride( + stride: Tuple[int], sizes: Optional[Tuple[int]] = None +) -> Tuple[bytes]: """ Dimension order represents how dimensions are laid out in memory, starting from the outer-most to the inner-most dimension. @@ -65,8 +67,16 @@ def dim_order_from_stride(stride: Tuple[int]) -> Tuple[bytes]: in original order. Thus when strides = (4, 3, 1, 1) returned value is (0, 1, 2, 3) Another example is: sizes = (1, 3, 1, 1) with strides = (3, 1, 3, 3), returned value is (0, 2, 3, 1) + + With sizes, ambiguous non-canonical orders are corrected when strides + exactly match channels-last. Preserve existing default orders: portable + kernels require matching dim orders even for physically equivalent layouts. """ - from torch.fx.experimental.symbolic_shapes import guard_or_false, guard_or_true + from torch.fx.experimental.symbolic_shapes import ( + guard_or_false, + guard_or_true, + statically_known_true, + ) for s in stride: torch._check(s != 0, lambda: "0 in strides is not supported for ExecuTorch.") @@ -93,6 +103,18 @@ def __lt__(self, other): sorted_dims = [ i[0] for i in sorted(enumerate(stride), key=lambda x: K(x[1]), reverse=True) ] + ndim = len(stride) + if ( + sizes is not None + and len(sizes) == ndim + and ndim in (4, 5) + and sorted_dims != list(range(ndim)) + ): + from torch._prims_common import make_channels_last_strides_for + + expected = make_channels_last_strides_for(sizes) + if all(statically_known_true(s == e) for s, e in zip(stride, expected)): + sorted_dims = [0, *range(2, ndim), 1] return tuple(typing.cast(Tuple[bytes], sorted_dims)) @@ -201,7 +223,7 @@ def from_tensor(cls, tensor: torch.Tensor, const: bool = False) -> TensorSpec: is_sparse=tensor.is_sparse, ) spec.stride = tensor.stride() - spec.dim_order = dim_order_from_stride(spec.stride) + spec.dim_order = dim_order_from_stride(spec.stride, tuple(spec.shape)) spec.requires_grad = tensor.requires_grad spec.storage = tensor.untyped_storage() if const else None diff --git a/exir/tensor_layout.py b/exir/tensor_layout.py index e902006783f..dfcdc9b3a74 100644 --- a/exir/tensor_layout.py +++ b/exir/tensor_layout.py @@ -35,5 +35,5 @@ def from_tensor(cls, tensor: torch.Tensor) -> "TensorLayout": return TensorLayout( scalar_type=scalar_type_enum(tensor.dtype), sizes=list(tensor.shape), - dim_order=list(dim_order_from_stride(tensor.stride())), + dim_order=list(dim_order_from_stride(tensor.stride(), tuple(tensor.shape))), ) diff --git a/exir/tests/test_dim_order_utils.py b/exir/tests/test_dim_order_utils.py index 6a8ea934db0..44a65025e19 100644 --- a/exir/tests/test_dim_order_utils.py +++ b/exir/tests/test_dim_order_utils.py @@ -7,9 +7,12 @@ # pyre-strict import unittest +from typing import List + import torch -from executorch.exir import to_edge_transform_and_lower +from executorch.exir import to_edge, to_edge_transform_and_lower from executorch.exir.dim_order_utils import get_dim_order, get_memory_format +from executorch.exir.schema import KernelCall, Tensor class TestDimOrderUtils(unittest.TestCase): @@ -47,3 +50,77 @@ def forward(self, t1, t2): expo_prog = torch.export.export(M, (x, y)) edge_prog = to_edge_transform_and_lower(expo_prog) edge_prog.to_executorch() + + def test_channels_last_single_channel_conv_dim_order(self) -> None: + class Conv(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(1, 8, kernel_size=3, padding=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x) + + model = Conv().to(memory_format=torch.channels_last).eval() + inputs = (torch.randn(1, 1, 12, 12).to(memory_format=torch.channels_last),) + edge_prog = to_edge(torch.export.export(model, inputs, strict=True)) + et_program = edge_prog.to_executorch().executorch_program + + conv_found = False + for inst in et_program.execution_plan[0].chains[0].instructions: + kernel = inst.instr_args + if not isinstance(kernel, KernelCall): + continue + op = et_program.execution_plan[0].operators[kernel.op_index] + if op.name != "aten::convolution": + continue + conv_found = True + tensors: List[Tensor] = [] + for arg in dict.fromkeys(kernel.args): + val = et_program.execution_plan[0].values[arg].val + if isinstance(val, Tensor) and len(val.sizes) == 4: + tensors.append(val) + # Input, weight, and output. + self.assertEqual(len(tensors), 3) + for tensor in tensors: + self.assertIn( + list(tensor.dim_order), + ([0, 1, 2, 3], [0, 2, 3, 1]), + ) + self.assertEqual(list(tensors[0].dim_order), list(tensors[-1].dim_order)) + self.assertTrue(conv_found) + + def test_singleton_dims_preserve_default_order(self) -> None: + class Add(torch.nn.Module): + def forward(self, x, y): + return x + y + + cases = ( + ( + Add(), + ( + torch.randn(2, 1, 3, 1).to(memory_format=torch.channels_last), + torch.randn(2, 1, 3, 1), + ), + ), + ( + torch.nn.ReLU(), + (torch.randn(2, 1, 3, 1, 1).to(memory_format=torch.channels_last_3d),), + ), + ) + for model, inputs in cases: + with self.subTest(model=type(model).__name__): + program = ( + to_edge_transform_and_lower(torch.export.export(model, inputs)) + .to_executorch() + .executorch_program + ) + tensors = [ + value.val + for value in program.execution_plan[0].values + if isinstance(value.val, Tensor) + ] + self.assertEqual(len(tensors), len(inputs) + 1) + for tensor in tensors: + self.assertEqual( + list(tensor.dim_order), list(range(len(tensor.sizes))) + ) diff --git a/exir/tests/test_tensor.py b/exir/tests/test_tensor.py index 19a536d44d5..d6307d716e1 100644 --- a/exir/tests/test_tensor.py +++ b/exir/tests/test_tensor.py @@ -321,6 +321,69 @@ def test_dim_order_from_stride_unbacked(self) -> None: with self.assertRaises(RuntimeError): dim_order_from_stride((u0, 0, 1)) + def test_dim_order_from_stride_with_sizes(self) -> None: + dim_order = dim_order_from_stride((490, 1, 10, 1), (1, 1, 49, 10)) + self.assertEqual((0, 2, 3, 1), dim_order) + self.assertEqual( + [490, 1, 10, 1], + stride_from_dim_order([1, 1, 49, 10], [0, 2, 3, 1]), + ) + + self.assertEqual( + (0, 1, 2, 3), + dim_order_from_stride((490, 490, 10, 1), (1, 1, 49, 10)), + ) + + self.assertEqual( + (0, 1, 2, 3), dim_order_from_stride((1, 1, 1, 1), (2, 1, 1, 1)) + ) + + self.assertEqual( + (0, 2, 3, 4, 1), + dim_order_from_stride((120, 1, 30, 6, 1), (1, 1, 4, 5, 6)), + ) + + self.assertEqual( + (3, 1, 2, 0), + dim_order_from_stride((1, 20, 5, 60), (2, 3, 4, 5)), + ) + + self.assertEqual((0, 2, 1, 3), dim_order_from_stride((490, 1, 10, 1))) + + t = torch.empty(1, 1, 49, 10).to(memory_format=torch.channels_last) + spec = TensorSpec.from_tensor(t) + self.assertEqual((0, 2, 3, 1), spec.dim_order) + + def test_dim_order_from_stride_preserves_supported_orders(self) -> None: + for sizes, strides in ( + ((2, 1, 3, 1), (3, 1, 1, 1)), + ((2, 1, 3, 1, 1), (3, 1, 1, 1, 1)), + ((2, 3, 4, 5), (60, 1, 15, 3)), + ((2, 3, 4, 5, 6), (360, 1, 90, 18, 3)), + ): + with self.subTest(sizes=sizes): + self.assertEqual( + dim_order_from_stride(strides), + dim_order_from_stride(strides, sizes), + ) + + def test_dim_order_from_stride_with_symbolic_sizes(self) -> None: + from torch.fx.experimental.symbolic_shapes import ShapeEnv + + shape_env = ShapeEnv() + height = shape_env.create_unbacked_symint() + torch._check_is_size(height) + torch._check(height >= 2) + self.assertEqual( + (0, 2, 3, 1), + dim_order_from_stride((10 * height, 1, 10, 1), (2, 1, height, 10)), + ) + self.assertEqual( + (0, 2, 1, 3), + dim_order_from_stride((490, 1, 10, 1), (2, 1, height, 10)), + ) + self.assertEqual(shape_env.guards, []) + def test_strides_from_dim_order(self) -> None: sizes = [] dim_order = []