executorch 1.4.0__cp314-cp314-win_amd64.whl

executorch/backends/arm/test/ops/test_maximum.py

@@ -0,0 +1,195 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+# Copyright 2024-2026 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.
+
+
+from typing import Tuple
+
+import torch
+from executorch.backends.arm.quantizer.arm_quantizer import (
+    get_symmetric_a16w8_quantization_config,
+)
+from executorch.backends.arm.test import common
+from executorch.backends.arm.test.tester.test_pipeline import (
+    EthosU55PipelineINT,
+    EthosU85PipelineINT,
+    TosaPipelineFP,
+    TosaPipelineINT,
+    VgfPipeline,
+)
+
+test_t = tuple[torch.Tensor, torch.Tensor]
+aten_op = "torch.ops.aten.maximum.default"
+exir_op = "executorch_exir_dialects_edge__ops_aten_maximum_default"
+
+
+class Maximum(torch.nn.Module):
+    test_parameters = {
+        "float_tensor": lambda: (
+            torch.FloatTensor([1, 2, 3, 5, 7]),
+            (torch.FloatTensor([2, 1, 2, 1, 10])),
+        ),
+        "ones": lambda: (torch.ones(1, 10, 4, 6), 2 * torch.ones(1, 10, 4, 6)),
+        "rand_diff": lambda: (torch.randn(1, 1, 4, 4), torch.ones(1, 1, 4, 1)),
+        "rand_same": lambda: (torch.randn(1, 3, 4, 4), torch.randn(1, 3, 4, 4)),
+        "rand_large": lambda: (
+            10000 * torch.randn(1, 1, 4, 4),
+            torch.randn(1, 1, 4, 1),
+        ),
+    }
+    test_parameters_fp16 = {
+        "fp16_rand": lambda: (
+            torch.randn(1, 3, 4, 4, dtype=torch.float16),
+            torch.randn(1, 3, 4, 4, dtype=torch.float16),
+        ),
+    }
+
+    test_parameters_bf16 = {
+        "bf16_rand": lambda: (
+            torch.randn(1, 3, 4, 4, dtype=torch.bfloat16),
+            torch.randn(1, 3, 4, 4, dtype=torch.bfloat16),
+        ),
+    }
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, y):
+        return torch.maximum(x, y)
+
+
+@common.parametrize(
+    "test_data",
+    Maximum.test_parameters
+    | Maximum.test_parameters_fp16
+    | Maximum.test_parameters_bf16,
+)
+def test_maximum_tosa_FP(test_data: Tuple):
+    TosaPipelineFP[test_t](
+        Maximum(), test_data(), aten_op, exir_op, tosa_extensions=["bf16"]
+    ).run()
+
+
+@common.parametrize("test_data", Maximum.test_parameters)
+def test_maximum_tosa_INT(test_data: Tuple):
+    TosaPipelineINT[test_t](Maximum(), test_data(), aten_op, exir_op).run()
+
+
+@common.parametrize("test_data", Maximum.test_parameters)
+@common.XfailIfNoCorstone300
+def test_maximum_u55_INT(test_data: Tuple):
+    EthosU55PipelineINT[test_t](
+        Maximum(),
+        test_data(),
+        aten_op,
+        exir_op,
+    ).run()
+
+
+@common.parametrize("test_data", Maximum.test_parameters)
+@common.XfailIfNoCorstone320
+def test_maximum_u85_INT(test_data: Tuple):
+    EthosU85PipelineINT[test_t](
+        Maximum(),
+        test_data(),
+        aten_op,
+        exir_op,
+    ).run()
+
+
+@common.parametrize(
+    "test_data",
+    Maximum.test_parameters
+    | Maximum.test_parameters_bf16
+    | Maximum.test_parameters_fp16,
+)
+@common.SkipIfNoModelConverter
+def test_maximum_vgf_no_quant(test_data: Tuple):
+    pipeline = VgfPipeline[test_t](
+        Maximum(),
+        test_data(),
+        aten_op,
+        exir_op,
+        quantize=False,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", Maximum.test_parameters)
+@common.SkipIfNoModelConverter
+def test_maximum_vgf_quant(test_data: Tuple):
+    pipeline = VgfPipeline[test_t](
+        Maximum(),
+        test_data(),
+        aten_op,
+        exir_op,
+        quantize=True,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", Maximum.test_parameters)
+@common.SkipIfNoModelConverter
+def test_maximum_vgf_quant_a16w8(test_data: Tuple):
+    pipeline = VgfPipeline[test_t](
+        Maximum(),
+        test_data(),
+        aten_op,
+        exir_op,
+        quantize=True,
+        tosa_extensions=["int16"],
+    )
+    pipeline.quantizer.set_global(get_symmetric_a16w8_quantization_config())
+    pipeline.run()
+
+
+@common.parametrize("test_data", Maximum.test_parameters)
+def test_maximum_tosa_INT_a16w8(test_data: test_t):
+    """Test maximum with 16A8W quantization for TOSA INT."""
+    pipeline = TosaPipelineINT[test_t](
+        Maximum(),
+        test_data(),
+        aten_op,
+        exir_op,
+        tosa_extensions=["int16"],
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", Maximum.test_parameters)
+@common.XfailIfNoCorstone300
+def test_maximum_u55_INT_a16w8(test_data: test_t):
+    """Test maximum with 16A8W quantization on U55 (16-bit activations, 8-bit
+    weights)
+    """
+    pipeline = EthosU55PipelineINT[test_t](
+        Maximum(),
+        test_data(),
+        aten_op,
+        exir_op,
+        per_channel_quantization=False,
+        a16w8_quantization=True,
+        use_to_edge_transform_and_lower=True,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", Maximum.test_parameters)
+@common.XfailIfNoCorstone320
+def test_maximum_u85_INT_a16w8(test_data: test_t):
+    """Test maximum with 16A8W quantization on U85 (16-bit activations, 8-bit
+    weights)
+    """
+    pipeline = EthosU85PipelineINT[test_t](
+        Maximum(),
+        test_data(),
+        aten_op,
+        exir_op,
+        per_channel_quantization=False,
+        a16w8_quantization=True,
+        use_to_edge_transform_and_lower=True,
+    )
+    pipeline.run()

executorch/backends/arm/test/ops/test_mean_dim.py

@@ -0,0 +1,487 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+# Copyright 2024-2026 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 pytest
+import torch
+from executorch.backends.arm.test import common
+from executorch.backends.arm.test.tester.test_pipeline import (
+    EthosU55PipelineINT,
+    EthosU85PipelineINT,
+    TosaPipelineFP,
+    TosaPipelineINT,
+    VgfPipeline,
+)
+
+input_t = tuple[torch.Tensor]
+
+
+class AdaptiveAveragePool2d(torch.nn.Module):
+    test_data_suite = {
+        # (test_name, test_data)
+        "zeros": lambda: (torch.zeros(1, 1280, 7, 7),),
+        "ones": lambda: (torch.ones(1, 1280, 7, 7),),
+        "rand": lambda: (torch.rand(1, 1280, 7, 7),),
+        "randn": lambda: (torch.randn(1, 1280, 7, 7),),
+    }
+    aten_op = "torch.ops.aten.adaptive_avg_pool2d.default"
+    exir_op = "executorch_exir_dialects_edge__ops_aten_mean_dim"
+
+    def __init__(self):
+        super().__init__()
+        self.adaptive_avg_pool2d = torch.nn.AdaptiveAvgPool2d(output_size=(1, 1))
+
+    def forward(self, x):
+        return self.adaptive_avg_pool2d(x)
+
+
+@common.parametrize("test_data", AdaptiveAveragePool2d.test_data_suite)
+def test_adaptive_avg_pool2d_tosa_FP(test_data):
+    TosaPipelineFP[input_t](
+        AdaptiveAveragePool2d(),
+        test_data(),
+        AdaptiveAveragePool2d.aten_op,
+        AdaptiveAveragePool2d.exir_op,
+    ).run()
+
+
+@common.parametrize("test_data", AdaptiveAveragePool2d.test_data_suite)
+def test_adaptive_avg_pool2d_tosa_INT(test_data):
+    TosaPipelineINT[input_t](
+        AdaptiveAveragePool2d(),
+        test_data(),
+        [],
+        [],
+        symmetric_io_quantization=True,
+    ).run()
+
+
+@common.parametrize("test_data", AdaptiveAveragePool2d.test_data_suite)
+@common.XfailIfNoCorstone300
+def test_adaptive_avg_pool2d_u55_INT(test_data):
+    EthosU55PipelineINT[input_t](
+        AdaptiveAveragePool2d(),
+        test_data(),
+        [],
+        [],
+        symmetric_io_quantization=True,
+    ).run()
+
+
+@common.parametrize("test_data", AdaptiveAveragePool2d.test_data_suite)
+@common.XfailIfNoCorstone320
+def test_adaptive_avg_pool2d_u85_INT(test_data):
+    EthosU85PipelineINT[input_t](
+        AdaptiveAveragePool2d(),
+        test_data(),
+        [],
+        [],
+        symmetric_io_quantization=True,
+    ).run()
+
+
+@common.parametrize("test_data", AdaptiveAveragePool2d.test_data_suite)
+@common.SkipIfNoModelConverter
+def test_adaptive_avg_pool2d_vgf_no_quant(test_data):
+    pipeline = VgfPipeline[input_t](
+        AdaptiveAveragePool2d(),
+        test_data(),
+        AdaptiveAveragePool2d.aten_op,
+        AdaptiveAveragePool2d.exir_op,
+        quantize=False,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", AdaptiveAveragePool2d.test_data_suite)
+@common.SkipIfNoModelConverter
+def test_adaptive_avg_pool2d_vgf_quant(test_data):
+    pipeline = VgfPipeline[input_t](
+        AdaptiveAveragePool2d(),
+        test_data(),
+        [],
+        [],
+        symmetric_io_quantization=True,
+        quantize=True,
+    )
+    pipeline.run()
+
+
+class MeanDim(torch.nn.Module):
+    test_data_suite: dict[str, tuple] = {
+        "rank_1_keepdim": lambda: (
+            torch.rand(7),
+            0,
+            True,
+        ),
+        "rank_2_keepdim": lambda: (
+            torch.rand(7, 3),
+            (0, 1),
+            True,
+        ),
+        "rank_3_keepdim": lambda: (
+            torch.rand(5, 7, 3),
+            (0, 1, 2),
+            True,
+        ),
+        "rand_1_keepdim": lambda: (
+            torch.rand(1, 5, 7, 3),
+            (1),
+            True,
+        ),
+        "rand_2_keepdim": lambda: (
+            torch.rand(1, 5, 7, 3),
+            (2),
+            True,
+        ),
+        "rand_3_keepdim": lambda: (
+            torch.rand(1, 5, 7, 3),
+            (3),
+            True,
+        ),
+        "rand_12_keepdim": lambda: (
+            torch.rand(1, 5, 7, 3),
+            (1, 2),
+            True,
+        ),
+        "rand_13_keepdim": lambda: (
+            torch.rand(1, 5, 7, 3),
+            (1, 3),
+            True,
+        ),
+        "rand_23_keepdim": lambda: (
+            torch.rand(1, 5, 7, 3),
+            (2, 3),
+            True,
+        ),
+        "rand_123_keepdim": lambda: (
+            torch.rand(1, 5, 7, 3),
+            (1, 2, 3),
+            True,
+        ),
+        "rand_0123_keepdim": lambda: (
+            torch.rand(1, 5, 7, 3),
+            (0, 1, 2, 3),
+            True,
+        ),
+        "rand_none_keepdim": lambda: (
+            torch.rand(1, 5, 7, 3),
+            None,
+            True,
+        ),
+        "rank_1": lambda: (
+            torch.rand(7),
+            (-1),
+            False,
+        ),
+        "rank_2": lambda: (
+            torch.rand(5, 7),
+            (-2, -1),
+            False,
+        ),
+        "rank_3": lambda: (
+            torch.rand(5, 7, 3),
+            (-3, -2, -1),
+            False,
+        ),
+        "rand_1": lambda: (
+            torch.rand(1, 5, 7, 3),
+            (-3),
+            False,
+        ),
+        "rand_2": lambda: (
+            torch.rand(1, 5, 7, 3),
+            (-2),
+            False,
+        ),
+        "rand_3": lambda: (
+            torch.rand(1, 5, 7, 3),
+            (-1),
+            False,
+        ),
+        "rand_12": lambda: (
+            torch.rand(1, 5, 7, 3),
+            (-3, -2),
+            False,
+        ),
+        "rand_13": lambda: (
+            torch.rand(1, 5, 7, 3),
+            (-3, -1),
+            False,
+        ),
+        "rand_23": lambda: (
+            torch.rand(1, 5, 7, 3),
+            (-2, -1),
+            False,
+        ),
+        "rand_123": lambda: (
+            torch.rand(1, 5, 7, 3),
+            (-3, -2, -1),
+            False,
+        ),
+        "rand_0123": lambda: (
+            torch.rand(1, 5, 7, 3),
+            (-4, -3, -2, -1),
+            False,
+        ),
+        "randn_large_axis": lambda: (
+            torch.randn(1, 256, 400),
+            (1, 2),
+            True,
+        ),
+        "rank5_01234": lambda: (
+            torch.rand(1, 1, 7, 3, 2),
+            (-5, -4, -3, -2, -1),
+            False,
+        ),
+        "rank5_234": lambda: (
+            torch.rand(1, 1, 7, 3, 2),
+            (-3, -2, -1),
+            False,
+        ),
+        "rank5_12": lambda: (
+            torch.rand(1, 1, 7, 3, 2),
+            (1, 2),
+            False,
+        ),
+        "rank5_2": lambda: (
+            torch.rand(1, 4, 7, 3, 2),
+            (2),
+            False,
+        ),
+    }
+    u55_test_data_suite = {
+        "u55_avg_pool_not_supported": lambda: (
+            torch.rand(1, 1, 1, 257),
+            (0, 1, 2, 3),
+            True,
+        ),
+    }
+    torch_op = "torch.ops.aten.mean.dim"
+    exir_op = "executorch_exir_dialects_edge__ops_aten_mean_dim"
+
+    def __init__(self, dim: int | list[int] = -1, keepdim: bool = True):
+        super().__init__()
+        self.dim = dim
+        self.keepdim = keepdim
+
+    def forward(self, x: torch.Tensor):
+        return x.mean(dim=self.dim, keepdim=self.keepdim)
+
+
+@common.parametrize("test_data", MeanDim.test_data_suite)
+def test_mean_dim_tosa_FP(test_data):
+    test_data, dim, keep_dim = test_data()
+    TosaPipelineFP[input_t](
+        MeanDim(dim, keep_dim),
+        (test_data,),
+        MeanDim.torch_op,
+        MeanDim.exir_op,
+    ).run()
+
+
+@common.parametrize("test_data", MeanDim.test_data_suite)
+def test_mean_dim_tosa_INT(test_data):
+    test_data, dim, keep_dim = test_data()
+    pipeline = TosaPipelineINT[input_t](
+        MeanDim(dim, keep_dim),
+        (test_data,),
+        [],  # Might be sum, avgpool, or both
+        symmetric_io_quantization=True,
+    )
+    pipeline.run()
+
+
+@common.parametrize(
+    "test_data", {**MeanDim.test_data_suite, **MeanDim.u55_test_data_suite}
+)
+@common.XfailIfNoCorstone300
+def test_mean_dim_u55_INT(test_data):
+    test_data, dim, keep_dim = test_data()
+    pipeline = EthosU55PipelineINT[input_t](
+        MeanDim(dim, keep_dim),
+        (test_data,),
+        [],  # Might be sum, avgpool, or both
+        symmetric_io_quantization=True,
+    )
+    pipeline.add_stage_after(
+        "export",
+        pipeline.tester.check_not,
+        ["torch.ops.aten.adaptive_avg_pool2d.default"],
+        suffix="avg_pool",
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", MeanDim.test_data_suite)
+@common.XfailIfNoCorstone320
+def test_mean_dim_u85_INT(test_data):
+    test_data, dim, keep_dim = test_data()
+    pipeline = EthosU85PipelineINT[input_t](
+        MeanDim(dim, keep_dim),
+        (test_data,),
+        [],  # Might be sum, avgpool, or both
+        symmetric_io_quantization=True,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", MeanDim.test_data_suite)
+@common.SkipIfNoModelConverter
+def test_mean_dim_vgf_no_quant(test_data):
+    test_data_val, dim, keep_dim = test_data()
+    pipeline = VgfPipeline[input_t](
+        MeanDim(dim, keep_dim),
+        (test_data_val,),
+        MeanDim.torch_op,
+        MeanDim.exir_op,
+        quantize=False,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", MeanDim.test_data_suite)
+@common.SkipIfNoModelConverter
+def test_mean_dim_vgf_quant(test_data):
+    test_data_val, dim, keep_dim = test_data()
+    pipeline = VgfPipeline[input_t](
+        MeanDim(dim, keep_dim),
+        (test_data_val,),
+        [],  # Quantized mean.dim may be decomposed to sum, avgpool, or both.
+        [],
+        symmetric_io_quantization=True,
+        quantize=True,
+    )
+    pipeline.run()
+
+
+class MeanDefault(torch.nn.Module):
+    torch_op = "torch.ops.aten.mean.default"
+    exir_op = "executorch_exir_dialects_edge__ops_aten_mean_default"
+
+    test_data_suite = {
+        "rank_1": lambda: (torch.rand(7),),
+        "rank_2": lambda: (torch.rand(2, 5),),
+        "rank_3": lambda: (torch.rand(2, 5, 7),),
+        "rank_4": lambda: (torch.rand(1, 5, 7, 3),),
+    }
+
+    def forward(self, x: torch.Tensor):
+        return torch.mean(x)
+
+
+@common.parametrize("test_data", MeanDefault.test_data_suite)
+def test_mean_tosa_FP(test_data):
+    pipeline = TosaPipelineFP[input_t](
+        MeanDefault(),
+        test_data(),
+        MeanDefault.torch_op,
+        MeanDefault.exir_op,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", MeanDefault.test_data_suite)
+def test_mean_tosa_INT(test_data):
+    pipeline = TosaPipelineINT[input_t](
+        MeanDefault(),
+        test_data(),
+        [],  # Might be sum, avgpool, or both
+        symmetric_io_quantization=True,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", MeanDefault.test_data_suite)
+@common.SkipIfNoModelConverter
+def test_mean_vgf_no_quant(test_data):
+    pipeline = VgfPipeline[input_t](
+        MeanDefault(),
+        test_data(),
+        MeanDefault.torch_op,
+        MeanDefault.exir_op,
+        quantize=False,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", MeanDefault.test_data_suite)
+@common.SkipIfNoModelConverter
+def test_mean_vgf_quant(test_data):
+    pipeline = VgfPipeline[input_t](
+        MeanDefault(),
+        test_data(),
+        [],  # Quantized mean.default may be decomposed to sum, avgpool, or both.
+        [],
+        symmetric_io_quantization=True,
+        quantize=True,
+    )
+    pipeline.run()
+
+
+a16w8_mean_test_parameters = {
+    "rank_2_keepdim": lambda: (torch.rand(7, 3), (0, 1), True),
+    "rank_3_keepdim": lambda: (torch.rand(5, 7, 3), (0, 1, 2), True),
+    "rand_23_keepdim": lambda: (torch.rand(1, 5, 7, 3), (2, 3), True),
+    "rand_0123_keepdim": lambda: (torch.rand(1, 5, 7, 3), (0, 1, 2, 3), True),
+    "rand_none_keepdim": lambda: (torch.rand(1, 5, 7, 3), None, True),
+    "rank_2": lambda: (torch.rand(5, 7), (-2, -1), False),
+    "rand_123": lambda: (torch.rand(1, 5, 7, 3), (-3, -2, -1), False),
+}
+
+a16w8_mean_test_parameters_u85_xfails = {
+    "rank_1_keepdim": lambda: (torch.rand(7), 0, True),
+    "rand_1_keepdim": lambda: (torch.rand(1, 5, 7, 3), (1), True),
+    "rank_1": lambda: (torch.rand(7), (-1), False),
+    "rand_3": lambda: (torch.rand(1, 5, 7, 3), (-1), False),
+}
+
+
+@common.parametrize(
+    "test_data", {**a16w8_mean_test_parameters, **a16w8_mean_test_parameters_u85_xfails}
+)
+@common.XfailIfNoCorstone300
+def test_mean_dim_a16w8_u55_INT(test_data):
+    test_data, dim, keep_dim = test_data()
+    pipeline = EthosU55PipelineINT[input_t](
+        MeanDim(dim, keep_dim),
+        (test_data,),
+        [],
+        symmetric_io_quantization=True,
+        a16w8_quantization=True,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", a16w8_mean_test_parameters)
+@common.XfailIfNoCorstone320
+def test_mean_dim_a16w8_u85_INT(test_data):
+    test_data, dim, keep_dim = test_data()
+    pipeline = EthosU85PipelineINT[input_t](
+        MeanDim(dim, keep_dim),
+        (test_data,),
+        [],
+        symmetric_io_quantization=True,
+        a16w8_quantization=True,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", a16w8_mean_test_parameters_u85_xfails)
+@pytest.mark.xfail(
+    reason="U85 mean_dim a16w8 produces all-zero output for single-dim reductions"
+)
+@common.XfailIfNoCorstone320
+def test_mean_dim_a16w8_u85_INT_xfail(test_data):
+    test_data, dim, keep_dim = test_data()
+    pipeline = EthosU85PipelineINT[input_t](
+        MeanDim(dim, keep_dim),
+        (test_data,),
+        [],
+        symmetric_io_quantization=True,
+        a16w8_quantization=True,
+    )
+    pipeline.run()