|
| 1 | +# Copyright 2024 Arm Limited and/or its affiliates. |
| 2 | +# All rights reserved. |
| 3 | +# |
| 4 | +# This source code is licensed under the BSD-style license found in the |
| 5 | +# LICENSE file in the root directory of this source tree. |
| 6 | + |
| 7 | +# |
| 8 | +# Tests the unsqueeze op which copies the data of the input tensor (possibly with new data format) |
| 9 | +# |
| 10 | + |
| 11 | +import unittest |
| 12 | +from typing import Sequence, Tuple |
| 13 | + |
| 14 | +import torch |
| 15 | + |
| 16 | +from executorch.backends.arm.quantizer.arm_quantizer import ( |
| 17 | + ArmQuantizer, |
| 18 | + get_symmetric_quantization_config, |
| 19 | +) |
| 20 | +from executorch.backends.arm.test import common |
| 21 | +from executorch.backends.arm.test.tester.arm_tester import ArmTester |
| 22 | + |
| 23 | +from executorch.backends.xnnpack.test.tester.tester import Quantize |
| 24 | +from parameterized import parameterized |
| 25 | + |
| 26 | + |
| 27 | +class TestSimpleUnsqueeze(unittest.TestCase): |
| 28 | + class Unsqueeze(torch.nn.Module): |
| 29 | + shapes: list[int | Sequence[int]] = [5, (5, 5), (5, 5), (5, 5, 5)] |
| 30 | + test_parameters: list[tuple[torch.Tensor]] = [(torch.ones(n),) for n in shapes] |
| 31 | + |
| 32 | + def forward(self, x: torch.Tensor, dim): |
| 33 | + return x.unsqueeze(dim) |
| 34 | + |
| 35 | + def _test_unsqueeze_tosa_MI_pipeline( |
| 36 | + self, module: torch.nn.Module, test_data: Tuple[torch.Tensor, int] |
| 37 | + ): |
| 38 | + ( |
| 39 | + ArmTester( |
| 40 | + module, |
| 41 | + example_inputs=test_data, |
| 42 | + compile_spec=common.get_tosa_compile_spec(), |
| 43 | + ) |
| 44 | + .export() |
| 45 | + .check_count({"torch.ops.aten.unsqueeze.default": 1}) |
| 46 | + .to_edge() |
| 47 | + .partition() |
| 48 | + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) |
| 49 | + .to_executorch() |
| 50 | + .run_method_and_compare_outputs(inputs=test_data) |
| 51 | + ) |
| 52 | + |
| 53 | + def _test_unsqueeze_tosa_BI_pipeline( |
| 54 | + self, module: torch.nn.Module, test_data: Tuple[torch.Tensor, int] |
| 55 | + ): |
| 56 | + quantizer = ArmQuantizer().set_io(get_symmetric_quantization_config()) |
| 57 | + ( |
| 58 | + ArmTester( |
| 59 | + module, |
| 60 | + example_inputs=test_data, |
| 61 | + compile_spec=common.get_tosa_compile_spec(), |
| 62 | + ) |
| 63 | + .quantize(Quantize(quantizer, get_symmetric_quantization_config())) |
| 64 | + .export() |
| 65 | + .check_count({"torch.ops.aten.unsqueeze.default": 1}) |
| 66 | + .to_edge() |
| 67 | + .partition() |
| 68 | + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) |
| 69 | + .to_executorch() |
| 70 | + .run_method_and_compare_outputs(inputs=test_data, qtol=1) |
| 71 | + ) |
| 72 | + |
| 73 | + def _test_unsqueeze_tosa_u55_pipeline( |
| 74 | + self, module: torch.nn.Module, test_data: Tuple[torch.Tensor, int] |
| 75 | + ): |
| 76 | + quantizer = ArmQuantizer().set_io(get_symmetric_quantization_config()) |
| 77 | + ( |
| 78 | + ArmTester( |
| 79 | + module, |
| 80 | + example_inputs=test_data, |
| 81 | + compile_spec=common.get_u55_compile_spec(), |
| 82 | + ) |
| 83 | + .quantize(Quantize(quantizer, get_symmetric_quantization_config())) |
| 84 | + .export() |
| 85 | + .check_count({"torch.ops.aten.unsqueeze.default": 1}) |
| 86 | + .to_edge() |
| 87 | + .partition() |
| 88 | + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) |
| 89 | + .to_executorch() |
| 90 | + ) |
| 91 | + |
| 92 | + @parameterized.expand(Unsqueeze.test_parameters) |
| 93 | + def test_unsqueeze_tosa_MI(self, test_tensor: torch.Tensor): |
| 94 | + for i in range(-test_tensor.dim() - 1, test_tensor.dim() + 1): |
| 95 | + self._test_unsqueeze_tosa_MI_pipeline(self.Unsqueeze(), (test_tensor, i)) |
| 96 | + |
| 97 | + @parameterized.expand(Unsqueeze.test_parameters) |
| 98 | + def test_unsqueeze_tosa_BI(self, test_tensor: torch.Tensor): |
| 99 | + self._test_unsqueeze_tosa_BI_pipeline(self.Unsqueeze(), (test_tensor, 0)) |
| 100 | + |
| 101 | + @parameterized.expand(Unsqueeze.test_parameters) |
| 102 | + def test_unsqueeze_u55_BI(self, test_tensor: torch.Tensor): |
| 103 | + self._test_unsqueeze_tosa_u55_pipeline(self.Unsqueeze(), (test_tensor, 0)) |
0 commit comments