test_linear.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
# pyre-unsafe

import unittest

from itertools import product
from typing import Callable, Dict, List, Optional, Tuple

import torch
from executorch.backends.xnnpack.partition.config.xnnpack_config import (
    ConfigPrecisionType,
)
from executorch.backends.xnnpack.partition.xnnpack_partitioner import (
    XnnpackFloatingPointPartitioner,
    XnnpackPartitioner,
)

from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import (
    get_symmetric_quantization_config,
)
from executorch.backends.xnnpack.quantizer.xnnpack_quantizer_utils import (
    QuantizationConfig,
)
from executorch.backends.xnnpack.test.tester import Quantize, Tester
from executorch.backends.xnnpack.test.tester.tester import (
    Partition,
    ToEdgeTransformAndLower,
)

from torch.export.graph_signature import ExportGraphSignature, InputKind

try:
    from torchao.quantization.quant_api import (
        int8_dynamic_activation_int4_weight,
        quantize_,
    )
    from torchao.utils import unwrap_tensor_subclass

    torchao_installed = True
except:
    torchao_installed = False


# Pytorch Modules Used for Testing
class BaseLinear(torch.nn.Module):
    def __init__(
        self,
        in_size: int = 2,
        input_channels: int = 4,
        output_channels: int = 4,
        dtype: torch.dtype = torch.float,
        use_bias: bool = False,
    ):
        super().__init__()
        self.linear = torch.nn.Linear(
            input_channels, output_channels, bias=use_bias
        ).to(dtype=dtype)

        self.ic = input_channels
        self.oc = output_channels

        assert dtype in [torch.float, torch.half], "Unsupported op dtype"
        self.op_dtype = dtype
        self.in_size = in_size

    def forward(self, x):
        return self.linear(x)

    def get_inputs(self, rank=3):
        # rank = 3 as default to inflate the act rank by 1 in batch dim
        # This is to make sure we don't specialize on 2D shapes.
        inp = torch.randn(self.in_size, self.ic).to(self.op_dtype)
        for _ in range(rank - 2):
            inp = inp.unsqueeze(0)
        assert inp.ndim == rank
        return (inp,)


class AddMMModule(torch.nn.Module):
    def __init__(self, in_size, out_size):
        super().__init__()
        self.mat = torch.nn.Parameter(torch.randn(in_size, out_size))
        self.bias = torch.nn.Parameter(torch.randn(1, out_size))

    def forward(self, x):
        return torch.addmm(self.bias, x, self.mat)


class LinearReluModule(torch.nn.Module):
    def __init__(self, in_size, out_size, use_bias, dtype=torch.float):
        super().__init__()
        self.dtype = dtype
        self.linear = torch.nn.Linear(in_size, out_size, bias=use_bias).to(dtype=dtype)

    def forward(self, x):
        return torch.nn.functional.relu(self.linear(x))

    def get_inputs(self):
        return (torch.randn(1, self.in_size, self.ic).to(self.op_dtype),)


class LinearParallelSequentialModule(torch.nn.Module):
    def __init__(
        self,
        in_size=2,
        input_size=4,
        intermediate_size=5,
        output_size=3,
        dtype=torch.float,
    ):
        super().__init__()
        self.linear1_weight = torch.nn.Parameter(
            torch.rand(intermediate_size, input_size)
        )
        self.linear1_bias = torch.nn.Parameter(torch.rand(intermediate_size))

        self.linear2_weight = torch.nn.Parameter(
            torch.rand(intermediate_size, input_size)
        )
        self.linear2_bias = torch.nn.Parameter(torch.rand(intermediate_size))

        self.linear3_weight = torch.nn.Parameter(
            torch.rand(output_size, intermediate_size)
        )
        self.linear3_bias = torch.nn.Parameter(torch.rand(output_size))
        self.in_size = in_size
        self.input_size = input_size
        self.dtype = torch.float

    def forward(self, x, y):
        a = torch.nn.functional.linear(x, self.linear1_weight, self.linear1_bias)
        b = torch.nn.functional.linear(y, self.linear2_weight, self.linear2_bias)
        c = torch.nn.functional.linear(b, self.linear3_weight, self.linear3_bias)
        return (a, c)

    def get_inputs(self):
        return (
            torch.rand(self.in_size, self.input_size, dtype=self.dtype),
            torch.rand(self.in_size, self.input_size, dtype=self.dtype),
        )


class LinearSequential(torch.nn.Module):
    def __init__(
        self,
        in_size=2,
        input_size=4,
        intermediate_size=5,
        output_size=3,
        dtype=torch.float,
    ):
        super().__init__()
        self.linear1_weight = torch.nn.Parameter(
            torch.rand(intermediate_size, input_size)
        )
        self.linear1_bias = torch.nn.Parameter(torch.rand(intermediate_size))

        self.linear2_weight = torch.nn.Parameter(
            torch.rand(output_size, intermediate_size)
        )
        self.linear2_bias = torch.nn.Parameter(torch.rand(output_size))
        self.in_size = in_size
        self.input_size = input_size
        self.dtype = torch.float

    def forward(self, x):
        a = torch.nn.functional.linear(x, self.linear1_weight, self.linear1_bias)
        b = torch.nn.functional.linear(a, self.linear2_weight, self.linear2_bias)
        return b

    def get_inputs(self):
        return (torch.rand(self.in_size, self.input_size, dtype=torch.float),)


class ParallelLinear(torch.nn.Module):
    def __init__(self, input_size, output_size):
        super().__init__()
        self.linear1_weight = torch.nn.Parameter(torch.rand(output_size, input_size))
        self.linear1_bias = torch.nn.Parameter(torch.rand(output_size))

        self.linear2_weight = torch.nn.Parameter(torch.rand(output_size, input_size))
        self.linear2_bias = torch.nn.Parameter(torch.rand(output_size))

    def forward(self, x, y):
        a = torch.nn.functional.linear(x, self.linear1_weight, self.linear1_bias)
        b = torch.nn.functional.linear(y, self.linear2_weight, self.linear2_bias)
        return a + b


class SharedDQChain(torch.nn.Module):
    def __init__(self, input_size, output_size):
        super().__init__()
        self.linear1_weight = torch.nn.Parameter(torch.rand(output_size, input_size))
        self.linear1_bias = torch.nn.Parameter(torch.rand(output_size))

        self.linear2_weight = torch.nn.Parameter(torch.rand(output_size, input_size))
        self.linear2_bias = torch.nn.Parameter(torch.rand(output_size))

    def forward(self, x):
        a = torch.nn.functional.linear(x, self.linear1_weight, self.linear1_bias)
        b = torch.nn.functional.linear(x, self.linear2_weight, self.linear2_bias)
        return a + b


class TestLinear(unittest.TestCase):
    """
    Test Class for XNNPACK Linear Operators.

    Notes:
        - XNNPACK Does not support Per Tensor Quantized Weights with Dynamic Activations
        - XNNPACK Only supports Per-Token Activation, so Dynamic per-tensor Quantization
          As done by the default dynamic quantization flow does Per-Token Quantization
          Activation under the hood, where the torch.nn.Module is doing Per-Tensor Quantization
          on the Activation. This is sufficient because Per-Token Quantization on Activations
          should produce strictly better results compared to Per-Tensor Quantization
    """

    def setUp(self):
        torch._dynamo.reset()

    @staticmethod
    def _get_4b_dqconfig() -> QuantizationConfig:
        # Returns a QuantizationConfig for 4b dynamic quantization for XNNPACK.
        qconfig: QuantizationConfig = get_symmetric_quantization_config(
            is_per_channel=True,
            is_dynamic=True,
            weight_qmin=-8,
            weight_qmax=7,
        )
        return qconfig

    def _test_linear(
        self,
        make_module,
        uses_bias,
        num_batch_dims=1,
        quant_type=None,
        dtype: torch.dtype = torch.float,
        atol=1e-03,  # TODO(T212995726): Investigate right atol for rand[n] inputs
    ):
        """
        Helper function to test linear op with different configurations.
        """
        edge_op = (
            "executorch_exir_dialects_edge__ops_aten_addmm_default"
            if uses_bias
            else "executorch_exir_dialects_edge__ops_aten_mm_default"
        )

        in_sizes = [3, 4, 4]
        input_sizes = [4, 37, 17]
        output_sizes = [4, 17, 37]

        quant_config = None
        if quant_type is not None:
            if quant_type == "per_channel":
                quant_config = get_symmetric_quantization_config(
                    is_per_channel=True,
                    is_dynamic=False,
                )
            elif quant_type == "per_tensor":
                quant_config = get_symmetric_quantization_config(
                    is_per_channel=False,
                    is_dynamic=False,
                )
            else:
                raise ValueError(f"Unsupported quant type {quant_type}")

        """
        Note that torch.nn.Linear maps to aten.mm.default (no bias) or aten.addmm.default (bias),
        which ares then transformed into aten.linear.default by the ConvertToLinear pass.
        """
        for i, _ in enumerate(in_sizes):
            torch._dynamo.reset()
            in_size = int(in_sizes[i])
            input_size = int(input_sizes[i])
            output_size = int(output_sizes[i])
            input_shape = [in_size] * num_batch_dims + [input_size]

            module = make_module(input_size, output_size).eval().to(dtype)
            inputs = (torch.randn(input_shape).to(dtype),)
            dynamic_shape = {}
            for i in range(num_batch_dims):
                dynamic_shape[i] = torch.export.Dim(f"batch{i}", min=2, max=in_size)

            dynamic_shape = (dynamic_shape,)

            for legacy_mode in (True, False):
                tester = Tester(module, inputs, dynamic_shapes=dynamic_shape)

                if quant_config:
                    tester.quantize(Quantize(quantization_config=quant_config))

                tester.export()
                if quant_config:
                    tester.check(["torch.ops.quantized_decomposed"])

                if legacy_mode:
                    tester.to_edge()
                    tester.partition()
                else:
                    tester.to_edge_transform_and_lower()

                tester.check_count(
                    {"torch.ops.higher_order.executorch_call_delegate": 1}
                )
                tester.check_not([edge_op])

                if quant_config:
                    tester.check_not(
                        [
                            "executorch_exir_dialects_edge__ops_aten_mm_default",
                            "executorch_exir_dialects_edge__ops_aten_addmm_default",
                        ]
                    )

                tester.to_executorch()
                tester.serialize()
                tester.run_method_and_compare_outputs(
                    qtol=bool(quant_config), atol=atol
                )

    def _test_dqlinear(
        self,
        module,
        inputs,
        dynamic_shapes,
        linear_count=1,
        is_per_channel=False,
        uses_bias=False,
        qconfig: Optional[QuantizationConfig] = None,
        atol=5e-02,  # TODO(T212995726): Investigate right atol for rand[n] inputs
    ):
        """
        Helper function to test dynamic quantized linear op with different configurations.
        """
        quant_config = qconfig or get_symmetric_quantization_config(
            is_per_channel=is_per_channel,
            is_dynamic=True,
        )
        for legacy_partitioner in (True, False):
            for per_op_mode in (True, False):
                DynamicallyQuantizedPartitioner = XnnpackPartitioner(
                    config_precisions=ConfigPrecisionType.DYNAMIC_QUANT,
                    per_op_mode=per_op_mode,
                )

                tester = Tester(module, inputs, dynamic_shapes=dynamic_shapes)
                tester.quantize(Quantize(quantization_config=quant_config))
                tester.export()

                if legacy_partitioner:
                    tester.to_edge()
                    tester.partition(Partition(DynamicallyQuantizedPartitioner))
                else:
                    tester.to_edge_transform_and_lower(
                        ToEdgeTransformAndLower([DynamicallyQuantizedPartitioner])
                    )
                tester.check_count(
                    {
                        "torch.ops.higher_order.executorch_call_delegate": (
                            linear_count if per_op_mode else 1
                        )
                    }
                )
                tester.check_not(
                    [
                        "executorch_exir_dialects_edge__ops_aten_mm_default",
                        "executorch_exir_dialects_edge__ops_aten_addmm_default",
                    ]
                )

                tester.to_executorch()
                tester.serialize()
                tester.run_method_and_compare_outputs(atol=atol)

    def _test_groupwise_dq_linear(
        self,
        mod: torch.nn.Module,
        inputs: Tuple[torch.Tensor],
        use_bias: bool = False,
        group_size: int = 8,
        num_linears: int = 1,
        atol: float = 5e-3,  # TODO(T212995726): Investigate right atol for rand[n] inputs
        rtol: float = 5e-3,  # TODO(T212995726): Investigate right rtol for rand[n] inputs
    ):
        """
        Helper function to test groupwise dynamic quantized linear op with different configurations.
        """
        quantize_(mod, int8_dynamic_activation_int4_weight(group_size=group_size))
        unwrap_tensor_subclass(mod)
        DynamicallyQuantizedPartitioner = XnnpackPartitioner(
            config_precisions=ConfigPrecisionType.DYNAMIC_QUANT,
            per_op_mode=True,
        )
        tester = (
            Tester(mod, inputs)
            .export()
            .check_count(
                {
                    "torch.ops.quant.choose_qparams_affine.default": 1 * num_linears,
                    "torch.ops.quant.quantize_affine.default": 1 * num_linears,
                    "torch.ops.quant.dequantize_affine.default": 2 * num_linears,
                    "torch.ops.aten.linear.default": 1 * num_linears,
                }
            )
        )
        (
            tester.to_edge_transform_and_lower(
                ToEdgeTransformAndLower([DynamicallyQuantizedPartitioner])
            )
        )

        (
            tester.check_count(
                {
                    "torch.ops.higher_order.executorch_call_delegate": 1,
                }
            )
            .check_not(
                [
                    "executorch_exir_dialects_edge__ops_quant_choose_qparams_affine_default",
                    "executorch_exir_dialects_edge__ops_quant_quantize_affine_default",
                    "executorch_exir_dialects_edge__ops_quant_dequantize_affine_default",
                    "executorch_exir_dialects_edge__ops_aten_mm_default",
                    "executorch_exir_dialects_edge__ops_aten_addmm_default",
                ]
            )
            .to_executorch()
            .serialize()
            .run_method_and_compare_outputs(atol=atol, rtol=rtol)
        )

    def _test_linear_overwrite_precision(
        self,
        make_module: Callable[[int, int], torch.nn.Module],
        uses_bias: bool,
        quant_type: str,
        quant_node_checks: List[Dict[str, int]],
        atol: float = 1e-03,  # TODO(T212995726): Investigate right atol for rand[n] inputs
    ):
        """
        This test is to test the overwrite precision of linear op.
        We will test partitioning, lowering, and running the quantized linear model as fp32 linear op.
        When using legacy_mode, we will test we don't partition [add]mm given,
        (1) We can't assume that weights are always static (non param).
        (2) Alternatively, when lowering [add]mm to xnn::bmm we can't support bias.
        (2)(a) Only lowering non-bias [add]mm, which is only exposed on legacy_path deemed low ROI.
        """

        in_sizes = [3, 4, 4]
        input_sizes = [4, 37, 17]
        output_sizes = [4, 17, 37]

        assert quant_type in ["per_tensor", "per_channel", "per_channel_dynamic"]
        per_channel = "per_channel" in quant_type
        dynamic = "dynamic" in quant_type
        quant_config = get_symmetric_quantization_config(
            is_per_channel=per_channel,
            is_dynamic=dynamic,
        )
        # Using FP32 partitioner for this quantized graph
        partitioner = XnnpackFloatingPointPartitioner()

        def get_qnode_checks(quant_node_checks, dialect):
            d = {}
            assert dialect in ["aten", "edge"]
            if dialect == "aten":
                d = {
                    f"torch.ops.quantized_decomposed.{op}": count
                    for op, count in quant_node_checks.items()
                }
            elif dialect == "edge":
                d = {
                    f"executorch.exir.dialects.edge._ops.quantized_decomposed.{op}".replace(
                        ".", "_"
                    ): count
                    for op, count in quant_node_checks.items()
                }
            assert len(d) == len(quant_node_checks)
            return d

        for i, _ in enumerate(in_sizes):
            torch._dynamo.reset()
            in_size = int(in_sizes[i])
            input_size = int(input_sizes[i])
            output_size = int(output_sizes[i])
            input_shape = [in_size] + [input_size]
            module = make_module(input_size, output_size).eval()
            inputs = (torch.randn(input_shape),)

            addmm_op_str = (
                "executorch_exir_dialects_edge__ops_aten_addmm_default"
                if uses_bias
                else "executorch_exir_dialects_edge__ops_aten_mm_default"
            )
            linear_op_str = "executorch_exir_dialects_edge__ops_aten_linear_default"

            for legacy_mode in (True, False):
                tester = (
                    Tester(module, inputs)
                    .quantize(Quantize(quantization_config=quant_config))
                    .export()
                    .dump_artifact()
                    .check_count(get_qnode_checks(quant_node_checks, "aten"))
                )

                if legacy_mode:
                    tester.to_edge()
                    tester.partition(Partition(partitioner=partitioner))
                    # We don't expect [add]mm to be partitioned
                    tester.check([addmm_op_str])
                else:
                    tester.to_edge_transform_and_lower(
                        ToEdgeTransformAndLower(partitioners=[partitioner])
                    )
                    # We do expect linear to be partitioned
                    tester.check_not([linear_op_str])

                # For legacy mode, fp32 permute_copy gets partitioned. (just a side effect)
                # For new mode, fp32 linear gets partitioned.
                tester.check_count(
                    {"torch.ops.higher_order.executorch_call_delegate": 1}
                )

                # Typically, we would not see any quantized ops in the graph.
                # But here we shouldn't partition these.
                tester.check_count(get_qnode_checks(quant_node_checks, "edge"))

                # TODO: Need to figure out how to load quantized ops in pybindings.
                # tester.to_executorch()
                # tester.serialize()
                # tester.run_method_and_compare_outputs(
                #     qtol=bool(quant_config), atol=atol
                # )

    def test_qd8_f32_per_channel_shared_dq_chain(self):
        for use_bias in (False, True):
            module = SharedDQChain(
                input_size=13,
                output_size=17,
            )
            inputs = (torch.randn(1, 2, 13),)

            self._test_dqlinear(
                module,
                inputs,
                dynamic_shapes=None,
                is_per_channel=True,
                linear_count=2,
                uses_bias=use_bias,
            )

    def _test_qd8_per_channel_linear(self, dtype: torch.dtype = torch.float):
        for uses_bias in (False, True):
            module = BaseLinear(
                in_size=8,
                input_channels=13,
                output_channels=17,
                dtype=dtype,
                use_bias=uses_bias,
            )
            inputs = module.get_inputs()

            self._test_dqlinear(
                module,
                inputs,
                dynamic_shapes=({1: torch.export.Dim("batch", max=100)},),
                is_per_channel=True,
                uses_bias=uses_bias,
            )

    def _test_qd8_linear_per_tensor_unsupported(self, dtype: torch.dtype = torch.float):
        for uses_bias in (False, True):
            module = BaseLinear(
                in_size=8,
                input_channels=13,
                output_channels=17,
                dtype=dtype,
                use_bias=uses_bias,
            )
            inputs = module.get_inputs()
            dynamic_shapes = ({1: torch.export.Dim("batch", max=100)},)

            quant_config = get_symmetric_quantization_config(
                is_per_channel=False,
                is_dynamic=True,
            )

            for legacy_partitioner in (True, False):
                for per_op_mode in (True, False):
                    # Every combination should fail to partition Linear or [add]mm.
                    DynamicallyQuantizedPartitioner = XnnpackPartitioner(
                        config_precisions=ConfigPrecisionType.DYNAMIC_QUANT,
                        per_op_mode=per_op_mode,
                    )

                    tester = Tester(module, inputs, dynamic_shapes=dynamic_shapes)
                    tester.quantize(Quantize(quantization_config=quant_config))
                    tester.export()

                    if legacy_partitioner:
                        tester.to_edge()
                        tester.partition(
                            Partition(DynamicallyQuantizedPartitioner)
                        ).dump_artifact()
                        # should have [add]mm node
                        if uses_bias:
                            tester.check(
                                [
                                    "executorch_exir_dialects_edge__ops_aten_addmm_default",
                                ]
                            )
                        else:
                            tester.check(
                                [
                                    "executorch_exir_dialects_edge__ops_aten_mm_default",
                                ]
                            )
                    else:
                        tester.to_edge_transform_and_lower(
                            ToEdgeTransformAndLower([DynamicallyQuantizedPartitioner])
                        ).dump_artifact()
                        # should not have a delegate node
                        tester.check_not(
                            [
                                "torch.ops.higher_order.executorch_call_delegate",
                            ]
                        )
                    # No need to run the model, since it should fail to partition.
                    return

    def _test_qd8_per_channel_4w_linear(self, dtype: torch.dtype = torch.float):
        qconfig = self._get_4b_dqconfig()
        input_channels = [2, 63]
        output_channels = [1, 127]
        batches = [
            2,
        ]
        use_bias = [False, True]
        dtypes = [
            dtype,
        ]

        for bs, bias, ipc, opc, dtype in product(
            batches,
            use_bias,
            input_channels,
            output_channels,
            dtypes,
        ):
            module = BaseLinear(
                in_size=bs,
                input_channels=ipc,
                output_channels=opc,
                dtype=dtype,
                use_bias=bias,
            )
            inputs = module.get_inputs()

            self._test_dqlinear(
                module,
                inputs,
                dynamic_shapes=({1: torch.export.Dim("batch", max=100)},),
                is_per_channel=True,
                uses_bias=bias,
                qconfig=qconfig,
                atol=5e-2,  # TODO(T212995726): Investigate right atol for rand[n] inputs
            )

    def _test_qd8_per_token_weight_per_channel_group_int4(
        self, dtype: torch.dtype = torch.float
    ):
        M_sizes = [1, 2, 17, 31]
        K_sizes = [32, 32, 64, 128]
        bl_sizes = [32, 32, 32, 64]
        N_sizes = [2, 17, 92, 128]

        for input_rank in range(2, 4):
            for use_bias in [True, False]:
                for M, K, bl, N in zip(M_sizes, K_sizes, bl_sizes, N_sizes):
                    lin_mod = BaseLinear(
                        in_size=M,
                        input_channels=K,
                        output_channels=N,
                        dtype=dtype,
                        use_bias=use_bias,
                    )

                    inputs = lin_mod.get_inputs(rank=input_rank)
                    # Half requires slightly higher atol, but if you look at error it is not that bad:
                    # Difference: max: 0.00140380859375, abs: 0.00140380859375, mean abs error: 0.00042724609375.
                    # -- Model vs. Reference --
                    # Numel: 4, 4
                    # Median: -0.05023193359375, -0.0516357421875
                    # Mean: 0.2373046875, 0.237060546875
                    # Max: 1.0078125, 1.0078125
                    # Min: -0.08465576171875, -0.08441162109375
                    atol = (
                        1e-2 if dtype == torch.half else 5e-3
                    )  # TODO(T212995726): Investigate right atol for rand[n] inputs
                    self._test_groupwise_dq_linear(
                        lin_mod, inputs, group_size=bl, use_bias=use_bias, atol=atol
                    )

    def test_fp16_linear(self):
        for use_bias in (True, False):
            for num_batch_dims in range(1, 3):
                self._test_linear(
                    lambda in_size, out_size: torch.nn.Linear(
                        in_size, out_size, bias=use_bias  # noqa
                    ),
                    num_batch_dims=num_batch_dims,
                    uses_bias=use_bias,
                    dtype=torch.float16,
                    atol=5e-2,  # TODO(T212995726): Investigate right atol for rand[n] inputs
                )

    def test_fp32_linear(self):
        for use_bias in (True, False):
            for num_batch_dims in range(1, 3):
                self._test_linear(
                    lambda in_size, out_size: torch.nn.Linear(
                        in_size, out_size, bias=use_bias  # noqa
                    ),
                    uses_bias=use_bias,
                    num_batch_dims=num_batch_dims,
                )

    def test_qc8_linear(self):
        for use_bias in (True, False):
            for num_batch_dims in range(1, 3):
                self._test_linear(
                    lambda in_size, out_size: torch.nn.Linear(
                        in_size, out_size, bias=use_bias  # noqa
                    ),
                    uses_bias=use_bias,
                    quant_type="per_channel",
                    num_batch_dims=num_batch_dims,
                )

    def test_fp32_addmm(self):
        # Note that the ConvertToLinear pass requires the weight matrix to be transposed.
        self._test_linear(
            lambda in_size, out_size: AddMMModule(in_size, out_size),
            uses_bias=True,
        )

    def test_fp32_linear_fused_relu(self):
        for use_bias in (True, False):
            for num_batch_dims in range(1, 3):
                self._test_linear(
                    lambda in_size, out_size: LinearReluModule(
                        in_size,
                        out_size,
                        use_bias,  # noqa
                    ),
                    uses_bias=use_bias,
                    num_batch_dims=num_batch_dims,
                )

    def test_qs8_linear_fused_relu(self):
        for use_bias in (True, False):
            for num_batch_dims in range(1, 3):
                self._test_linear(
                    lambda in_size, out_size: LinearReluModule(
                        in_size,
                        out_size,
                        use_bias,  # noqa
                    ),
                    num_batch_dims=num_batch_dims,
                    uses_bias=use_bias,
                    quant_type="per_tensor",
                )

    def test_qs8_linear(self):
        for use_bias in (True, False):
            for num_batch_dims in range(1, 3):
                self._test_linear(
                    lambda in_size, out_size: torch.nn.Linear(
                        in_size, out_size, bias=use_bias  # noqa
                    ),
                    uses_bias=use_bias,
                    num_batch_dims=num_batch_dims,
                    quant_type="per_tensor",
                )

    # Tests for q[dp]8-f16-qc8w
    def test_qd8_f16_per_channel_linear(self):
        self._test_qd8_per_channel_linear(dtype=torch.half)

    def test_qd8_f16_per_tensor_linear(self):
        """
        XNNPACK doesn't support per_tensor quantized weights for dynamic quantized linear op.
        This test is to verify that we can't lower per_tensor quantized weights to per_channel quantized weights.
        """
        self._test_qd8_linear_per_tensor_unsupported(dtype=torch.half)

    # Tests for q[dp]8-f32-qc8w
    def test_qd8_f32_per_channel_linear(self):
        self._test_qd8_per_channel_linear(dtype=torch.float)

    def test_qd8_f32_per_tensor_linear(self):
        """
        XNNPACK doesn't support per_tensor quantized weights for dynamic quantized linear op.
        This test is to verify that we can't lower per_tensor quantized weights to per_channel quantized weights.
        """
        self._test_qd8_linear_per_tensor_unsupported(dtype=torch.half)

    # Tests for q[dp]8-f16-qc4w
    def test_linear_qd8_f16_per_channel_int4(self):
        self._test_qd8_per_channel_4w_linear(dtype=torch.half)

    # Tests for q[dp]8-f32-qc4w
    def test_linear_qd8_f32_per_channel_int4(self):
        self._test_qd8_per_channel_4w_linear(dtype=torch.float)

    # Tests for q[dp]8-f16-qb4w
    @unittest.skipIf(
        not torchao_installed, "Per Channel Group Quantization Required TorchAO"
    )
    def test_linear_qd8_f16_per_token_weight_per_channel_group_int4(self):
        self._test_qd8_per_token_weight_per_channel_group_int4(dtype=torch.half)

    # Tests for q[dp]8-f32-qb4w
    @unittest.skipIf(
        not torchao_installed, "Per Channel Group Quantization Required TorchAO"
    )
    def test_linear_qd8_f32_per_token_weight_per_channel_group_int4(self):
        self._test_qd8_per_token_weight_per_channel_group_int4(dtype=torch.float)

    @unittest.skipIf(
        not torchao_installed, "Per Channel Group Quantization Required TorchAO"
    )
    def test_linear_qd8_per_token_groupwise_unsupported_groupsize(self):
        # groupsize must be multiple of 32
        for dtype in [torch.float, torch.half]:
            lin_mod = BaseLinear(
                in_size=1,
                input_channels=60,
                output_channels=60,
                dtype=dtype,
                use_bias=True,
            )
            inputs = lin_mod.get_inputs()

            with self.assertRaisesRegex(
                AssertionError,
                "Delegation to XNNPACK requires group_size to be a multiple of 32, but got 30",
            ):
                self._test_groupwise_dq_linear(
                    lin_mod, inputs, group_size=30, use_bias=False, atol=1e-2
                )

    def test_qd8_per_channel_linear_parallel(self):
        in_size = 2
        input_size = 4
        output_size = 5
        for dtype in torch.float, torch.half:
            inputs = (
                torch.rand(in_size, input_size, dtype=dtype),
                torch.rand(in_size, input_size, dtype=dtype),
            )
            batch_dim = torch.export.Dim("batch", max=100)
            dynamic_shapes = ({0: batch_dim}, {0: batch_dim})

            self._test_dqlinear(
                ParallelLinear(input_size=input_size, output_size=output_size).to(
                    dtype
                ),
                inputs,
                dynamic_shapes=dynamic_shapes,
                linear_count=2,
                is_per_channel=True,
                uses_bias=True,
            )

    def test_qd8_per_channel_linear_with_two_batch(self):
        in_size = 2
        input_size = 14
        output_size = 15

        for dtype in torch.float, torch.half:
            for use_bias in (False, True):
                linear = BaseLinear(
                    in_size=in_size,
                    input_channels=input_size,
                    output_channels=output_size,
                    dtype=dtype,
                    use_bias=use_bias,
                )
                # Create inputs with two batch dimensions, i.e. 3D activation
                inputs = (torch.randn(in_size, in_size, input_size).to(dtype),)
                batch_dim = torch.export.Dim("batch", max=100)
                dynamic_shapes = ({0: batch_dim, 1: batch_dim},)

                self._test_dqlinear(
                    linear,
                    inputs,
                    dynamic_shapes=dynamic_shapes,
                    linear_count=1,
                    is_per_channel=True,
                    uses_bias=True,
                )

    def test_qd8_per_channel_linear_sequential(self):
        lin_mod = LinearSequential()
        inputs = lin_mod.get_inputs()
        dynamic_shapes = ({0: torch.export.Dim("batch", max=100)},)

        self._test_dqlinear(
            lin_mod,
            inputs,
            dynamic_shapes=dynamic_shapes,
            linear_count=2,
            is_per_channel=True,
            uses_bias=True,
            atol=1e-1,  # TODO(T212995726): Investigate right atol for rand[n] inputs
        )

    def test_qd8_per_channel_linear_parallel_and_sequential(self):
        lin_mod = LinearParallelSequentialModule()
        inputs = lin_mod.get_inputs()
        dynamic_shapes = (
            {0: torch.export.Dim("batch", max=100)},
            {0: torch.export.Dim("batch2", max=100)},
        )

        self._test_dqlinear(
            lin_mod,
            inputs,
            dynamic_shapes=dynamic_shapes,
            linear_count=3,
            is_per_channel=True,
            uses_bias=True,
            atol=1e-1,  # TODO(T212995726): Investigate right atol for rand[n] inputs
        )

    def test_linear_qs8_as_fp32(self):
        for use_bias in (True, False):
            self._test_linear_overwrite_precision(
                lambda in_size, out_size: torch.nn.Linear(
                    in_size, out_size, bias=use_bias  # noqa
                ),
                use_bias,
                "per_tensor",
                quant_node_checks={
                    "quantize_per_tensor.default": 2,  # 1: act, 1: output
                    "dequantize_per_tensor.default": 3,  # 1: act, 1: weight, 1: output
                },
            )

    def test_linear_qc8_as_fp32(self):
        for use_bias in (True, False):
            self._test_linear_overwrite_precision(
                lambda in_size, out_size: torch.nn.Linear(
                    in_size, out_size, bias=use_bias  # noqa
                ),
                use_bias,
                "per_channel",
                quant_node_checks={
                    "quantize_per_tensor.default": 2,  # 1: act, 1: output
                    "dequantize_per_tensor.default": 2,  # 1: act, 1: output
                    "dequantize_per_channel.default": 1,  # 1: weight
                },
            )

    def test_linear_qd8_as_fp32(self):
        for use_bias in (True, False):
            self._test_linear_overwrite_precision(
                lambda in_size, out_size: torch.nn.Linear(
                    in_size, out_size, bias=use_bias  # noqa
                ),
                use_bias,
                "per_channel_dynamic",
                quant_node_checks={
                    "quantize_per_tensor.tensor": 1,  # 1: act
                    "dequantize_per_tensor.tensor": 1,  #  1: act
                    "dequantize_per_channel.default": 1,  # 1: weight
                },
            )

    def test_linear_with_force_non_static_weights_for_f32_linear(self):
        def check_signature(
            signature: ExportGraphSignature,
            force_flag: bool,
            use_bias: bool,
            legacy_mode: bool,
        ):
            num_params = 0
            if force_flag:
                num_params = 1  # weight_param
                if use_bias:
                    num_params += 1  # bias_param
            sign_params: int = 0
            input_specs = signature.input_specs
            for input_spec in input_specs:
                if input_spec.kind == InputKind.PARAMETER:
                    sign_params += 1
            assert (
                sign_params == num_params
            ), f"Expected {num_params} params, got {sign_params} with force_flag={force_flag}, use_bias={use_bias}, legacy_mode={legacy_mode}"

        for force_flag in (True, False):
            for use_bias in (True, False):
                for legacy_mode in (True, False):
                    module = BaseLinear(
                        in_size=8,
                        input_channels=13,
                        output_channels=17,
                        use_bias=use_bias,
                    )
                    inputs = module.get_inputs()
                    tester = Tester(module, inputs).export()
                    partitioner = XnnpackPartitioner(
                        force_non_static_weights_for_f32_linear=force_flag
                    )
                    if legacy_mode:
                        tester.to_edge()
                        partitioner_stage = Partition(partitioner=partitioner)
                        tester.partition(partition_stage=partitioner_stage)
                        tester.check_not(
                            [
                                (
                                    "executorch_exir_dialects_edge__ops_aten_mm_default"
                                    if use_bias
                                    else "executorch_exir_dialects_edge__ops_aten_addmm_default"
                                )
                            ]
                        )
                    else:
                        to_edge_and_transform_stage = ToEdgeTransformAndLower(
                            partitioners=[partitioner]
                        )
                        tester.to_edge_transform_and_lower(
                            to_edge_and_transform_stage=to_edge_and_transform_stage
                        )
                        tester.check_not(
                            ["executorch_exir_dialects_edge__ops_aten_linear_default"]
                        )

                    signature: ExportGraphSignature = (
                        tester.get_artifact().exported_program().graph_signature
                    )
                    check_signature(signature, force_flag, use_bias, legacy_mode)

                    tester.to_executorch()
                    tester.serialize()
                    tester.run_method_and_compare_outputs()