quantized_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.

import operator

from typing import Optional

import executorch.backends.vulkan.utils as utils

import torch
import torch.nn.functional as F

from executorch.backends.transforms.utils import (
    create_constant_placeholder,
    get_param_tensor,
)

from executorch.backends.vulkan.patterns.pattern_registry import (
    PatternMatch,
    register_pattern_detector,
    register_pattern_replacement,
)

from executorch.exir import ExportedProgram
from executorch.exir.dialects._ops import ops as exir_ops

from torch.export.graph_signature import InputKind


class QuantizedLinearMatch(PatternMatch):
    def __init__(self, mm_node: torch.fx.Node) -> None:  # noqa: C901
        self.anchor_node = mm_node
        self.match_found = False
        self.all_nodes = [self.anchor_node]

        # addmm(bias, mat1, mat2) has a different arg layout than
        # mm(mat1, mat2) and linear(input, weight, bias?)
        is_addmm = self.anchor_node.target == exir_ops.edge.aten.addmm.default
        weight_arg_idx = 2 if is_addmm else 1
        input_arg_idx = 1 if is_addmm else 0

        const_node, arg_chain = utils.trace_args_until_placeholder(
            self.anchor_node.args[weight_arg_idx]
        )

        # mat2 is not a constant tensor - no match
        if const_node is None:
            return

        dequantize_weight_node = None
        # Search for a dequantize node in the arg chain of weight
        for node in arg_chain:
            if isinstance(node, torch.fx.Node) and utils.is_dequant_node(node):
                dequantize_weight_node = node
        # weight is not quantized - no match
        if dequantize_weight_node is None:
            return

        self.weight_node = const_node
        self.dequantize_weight_node = dequantize_weight_node
        self.all_nodes.extend(arg_chain)

        # By default, assume dequant node is from quantized_decomposed namespace
        scales_arg_idx = 1
        zeros_arg_idx = 2
        # torchao dequantize has a different function schema than quantized_decomposed
        if (
            self.dequantize_weight_node.target
            == exir_ops.edge.torchao.dequantize_affine.default
        ):
            scales_arg_idx = 2
            zeros_arg_idx = 3

        # Identify weight quantization parameter nodes
        self.weight_scales_node, arg_chain = utils.trace_args_until_placeholder(
            self.dequantize_weight_node.args[scales_arg_idx]
        )
        assert self.weight_scales_node is not None
        self.all_nodes.extend(arg_chain)

        self.weight_zeros_node, arg_chain = utils.trace_args_until_placeholder(
            self.dequantize_weight_node.args[zeros_arg_idx]
        )
        assert self.weight_zeros_node is not None
        self.all_nodes.extend(arg_chain)

        # Identify output node
        self.output_node = self.anchor_node

        # Identify primary input node of the anchor. Due to decomposition of aten.linear
        # there may be a view_copy node between the original input tensor to the linear
        # op and the actual linear op node.
        anchor_primary_input_node = self.anchor_node.args[input_arg_idx]
        assert isinstance(anchor_primary_input_node, torch.fx.Node)

        # Skip potential view_copy between dq and linear
        if utils.is_view_copy_node(anchor_primary_input_node):
            self.all_nodes.append(anchor_primary_input_node)
            anchor_primary_input_node = anchor_primary_input_node.args[
                0
            ]  # pyre-ignore[16]
            assert isinstance(anchor_primary_input_node, torch.fx.Node)

        # By default, assume that the input tensor is not quantized in any way
        self.quantize_input_node = None
        self.dequantize_input_node = None
        self.pattern_input_node = anchor_primary_input_node

        self.input_scales_node = None
        self.input_zeros_node = None

        scales_arg_idx = 1
        zeros_arg_idx = 2

        # If the primary input node comes from a dequantize node, that implies the input
        # input tensor is quantized (either statically or dynamically).
        if utils.is_dequant_node(anchor_primary_input_node):
            # Assume that this is a static quantization pattern; the input to the
            # pattern is a statically quantized int8 tensor.
            self.dequantize_input_node = anchor_primary_input_node
            self.all_nodes.append(self.dequantize_input_node)
            input_to_dq_node = self.dequantize_input_node.args[0]
            self.pattern_input_node = input_to_dq_node

            # torchao dequantize has a slightly different function schema
            if (
                self.dequantize_input_node.target
                == exir_ops.edge.torchao.dequantize_affine.default
            ):
                scales_arg_idx = 2
                zeros_arg_idx = 3

            self.input_scales_node = self.dequantize_input_node.args[scales_arg_idx]
            self.input_zeros_node = self.dequantize_input_node.args[zeros_arg_idx]

            # Check for dynamic quantization: input scales are dynamically
            # computed via a choose_qparams op
            if utils.is_quant_node(input_to_dq_node) and utils.is_dynamic_qscale(
                self.input_scales_node
            ):
                self.quantize_input_node = input_to_dq_node
                self.pattern_input_node = self.quantize_input_node.args[0]

        # The implementation has a limitation that input channels must be a
        # multiple of 4. This is to ensure that data loads are aligned well with
        # texel boundaries. If this is not true, then don't match the pattern.
        in_channels = self.pattern_input_node.meta["val"].shape[-1]
        if in_channels % 4 != 0:
            return

        # Identify bias node, if applicable
        self.bias_node = None
        if self.anchor_node.target == exir_ops.edge.aten.addmm.default:
            self.bias_node, arg_chain = utils.trace_args_until_placeholder(
                self.anchor_node.args[0]
            )
            assert self.bias_node is not None
            self.all_nodes.extend(arg_chain)
        elif self.anchor_node.target == exir_ops.edge.aten.linear.default:
            if len(self.anchor_node.args) > 2 and self.anchor_node.args[2] is not None:
                self.bias_node, arg_chain = utils.trace_args_until_placeholder(
                    self.anchor_node.args[2]
                )
                if self.bias_node is not None:
                    self.all_nodes.extend(arg_chain)

        # If input is not quantized, then we are done
        if self.dequantize_input_node is None:
            self.match_found = True
            return

        # Check if the output is also quantized (q → dq → linear → q pattern)
        # Also handle fused linear+relu (q → dq → linear → relu → q pattern)
        # Due to decomposition of aten.linear for 3D+ inputs, there may be a
        # view_copy between the mm output and the quantize node.
        self.quantize_output_node = None
        self.output_scales_node = None
        self.output_zeros_node = None
        self.relu_node = None
        self.output_view_copy_node = None
        if len(self.output_node.users) == 1:
            cur_node = list(self.output_node.users)[0]
            # Skip potential view_copy between linear and output quantize
            if utils.is_view_copy_node(cur_node) and len(cur_node.users) == 1:
                self.output_view_copy_node = cur_node
                self.all_nodes.append(self.output_view_copy_node)
                self.output_node = self.output_view_copy_node
                cur_node = list(cur_node.users)[0]
            if cur_node.target == exir_ops.edge.aten.relu.default:
                self.relu_node = cur_node
                if len(cur_node.users) == 1:
                    cur_node = list(cur_node.users)[0]
                else:
                    cur_node = None
            if cur_node is not None and utils.is_quant_node(cur_node):
                self.quantize_output_node = cur_node
                self.output_scales_node = self.quantize_output_node.args[1]
                self.output_zeros_node = self.quantize_output_node.args[2]

        self.match_found = True

    def is_weight_only_quantized(self) -> bool:
        return self.dequantize_input_node is None

    def has_output_quantization(self) -> bool:
        return (
            hasattr(self, "quantize_output_node")
            and self.quantize_output_node is not None
        )

    def is_weight_pergroup_quantized(self) -> bool:
        weight_shape = self.weight_node.meta["val"].shape
        scales_shape = self.weight_scales_node.meta["val"].shape
        if len(scales_shape) != 2:
            return False

        # Check that:
        # height dim of scales is same as height dim of weight (N / output channels dim)
        # width dim of weight (K / in channels dim) is divisible by width dim of scales
        # (number of quantization groups)
        return scales_shape[-2] == weight_shape[-2] and (
            weight_shape[-1] % scales_shape[-1] == 0
        )

    def is_weight_perchannel_quantized(self) -> bool:
        weight_shape = self.weight_node.meta["val"].shape
        scales_shape = self.weight_scales_node.meta["val"].shape
        if len(scales_shape) != 1:
            return False

        # scales should have same size as weight's output channels dim
        return scales_shape[0] == weight_shape[-2]

    def is_input_static_per_tensor_quantized(self) -> bool:
        if self.dequantize_input_node is None:
            return False

        # For static quantization per tensor quantization, the scales and zeros
        # are scalars.
        return isinstance(self.input_scales_node, float)

    def is_input_dynamic_perchannel_quantized(self) -> bool:
        if self.dequantize_input_node is None:
            return False

        if not isinstance(self.input_scales_node, torch.fx.Node):
            return False

        # For dynamic quantization, input scale node should be a getitem operator
        # retrieving the output of a choose_qparams op
        if self.input_scales_node.target != operator.getitem:
            return False

        # The getitem node should be retrieving from a choose_qparams op
        if not utils.is_choose_qparams_node(self.input_scales_node.args[0]):
            return False

        scales_shape = self.input_scales_node.meta["val"].shape
        input_shape = self.pattern_input_node.meta["val"].shape

        return input_shape[-2] == scales_shape[-1]


linear_anchor_nodes = {
    exir_ops.edge.aten.linear.default,
    exir_ops.edge.aten.mm.default,
    exir_ops.edge.aten.addmm.default,
}


@register_pattern_detector("quantized_linear")
def find_quantized_linear_patterns(
    node: torch.fx.Node,
) -> Optional[QuantizedLinearMatch]:
    if node.target not in linear_anchor_nodes:
        return None

    matched_pattern = QuantizedLinearMatch(node)
    if matched_pattern.match_found:
        return matched_pattern

    return None


##
## Constant tensor manipulation
##


def pack_4bit_weight_tensor(weight_tensor: torch.Tensor) -> torch.Tensor:
    """
    Given a 8-bit weight tensor containing values quantized to 4 bits, create a packed
    weight tensor by transposing the weight tensor, then packing 2 4-bit values in one
    8-bit value.

    An input weight tensor of shape (N, K) will produce a packed weight tensor of shape
    (K, N / 2).
    """

    # Assert we got a properly quantized tensor.
    min_val, max_val = weight_tensor.min().item(), weight_tensor.max().item()
    assert (
        max_val <= 7 and min_val >= -8
    ), f"pack_4bit_weight_tensor: [min_val,max_val] out of [-8, 7] range, got [{min_val}, {max_val}]"

    # Assuming we have a 2d tensor
    if weight_tensor.ndim != 2:
        weight_tensor = weight_tensor.squeeze()
    assert (
        weight_tensor.ndim == 2
    ), f"pack_4bit_weight_tensor: expecting input tensor to be 2d, got {weight_tensor.ndim}"

    # Need to pad innermost dim to be a multiple of 8, since the minimum load granularity
    # is int32 (4 bytes), which contains 8 4-bit values.
    if weight_tensor.shape[-1] % 8 != 0:
        num_pad = 8 - (weight_tensor.shape[-1] % 8)
        weight_tensor = F.pad(input=weight_tensor, pad=(0, num_pad))

    # Shape after padding
    _, in_channels = weight_tensor.shape
    assert in_channels % 8 == 0, "convert_to_qc4w: expecting ic to be divisible by 8"

    # Adjust weight_tensor tensor for zp
    weight_tensor = weight_tensor.to(dtype=torch.uint8) + 8
    # Pack each 4-bit value into a single 8-bit value
    return weight_tensor[::, 1::2] << 4 | weight_tensor[::, ::2]


def compute_per_group_sums(weight_tensor: torch.Tensor, group_size: int):
    """
    Compute the sum of weights per quantization group.

    Args:
        weight_tensor (torch.Tensor): Tensor of shape [out_channels, in_channels], dtype int8.
        group_size (int): Number of input channels per quantization group.

    Returns:
        torch.Tensor: Tensor of shape [num_groups, out_channels], where num_groups = in_channels // group_size.
    """
    out_channels, in_channels = weight_tensor.shape
    num_groups = in_channels // group_size
    # Reshape to [out_channels, num_groups, group_size]
    reshaped = weight_tensor.view(out_channels, num_groups, group_size)
    # Sum over group_size dimension to get [out_channels, num_groups]
    sums = reshaped.sum(dim=2)
    # Transpose to [num_groups, out_channels]
    sums = sums.transpose(0, 1).contiguous()
    # Pad out_channels dim (dim=1) to be a multiple of 8 if needed
    out_channels = sums.shape[1]
    if out_channels % 8 != 0:
        num_pad = 8 - (out_channels % 8)
        sums = F.pad(sums, (0, num_pad))

    return sums.to(torch.int32).contiguous()


##
## Pattern Replacement
##


def make_linear_q4gsw_op(
    ep: ExportedProgram,
    graph_module: torch.fx.GraphModule,
    match: QuantizedLinearMatch,
    weight_tensor: torch.Tensor,
    weight_scales_tensor: torch.Tensor,
):
    num_groups = weight_scales_tensor.shape[-1]
    in_channels = weight_tensor.shape[-1]
    group_size = in_channels // num_groups

    weight_tensor = pack_4bit_weight_tensor(weight_tensor)
    # Use this function for convenience to update the state dict with the packed
    # weight tensor. Alignment will already have been done in the above function.
    weight_tensor = utils.align_width_and_update_state_dict(
        ep, match.weight_node, weight_tensor, align_to=1, force_update=True
    )

    # Also transpose the weight scales tensor to shape [num_groups, N]
    weight_scales_tensor = weight_scales_tensor.transpose(0, 1).contiguous()
    # Align to multiple of 8 to ensure that data loads from the weight scales
    # tensor do not go out of bounds. Each thread computes 8 output channels.
    utils.align_width_and_update_state_dict(
        ep,
        match.weight_scales_node,
        weight_scales_tensor,
        align_to=8,
        force_update=True,
    )

    with graph_module.graph.inserting_before(match.output_node):
        linear_q4gsw_node = graph_module.graph.create_node(
            "call_function",
            exir_ops.edge.et_vk.linear_q4gsw.default,
            args=(
                match.pattern_input_node,
                match.weight_node,
                match.weight_scales_node,
                group_size,
            ),
        )

    linear_q4gsw_node.meta["val"] = match.output_node.meta["val"]
    match.output_node.replace_all_uses_with(linear_q4gsw_node)


def make_linear_dq8ca_q4gsw_op(
    ep: ExportedProgram,
    graph_module: torch.fx.GraphModule,
    match: QuantizedLinearMatch,
    weight_tensor: torch.Tensor,
    weight_scales_tensor: torch.Tensor,
):
    num_groups = weight_scales_tensor.shape[-1]
    in_channels = weight_tensor.shape[-1]
    group_size = in_channels // num_groups

    # Compute per quant group sums before packing the weight tensor
    sum_per_quant_group = compute_per_group_sums(weight_tensor, group_size)

    weight_tensor = pack_4bit_weight_tensor(weight_tensor)
    # Use this function for convenience to update the state dict with the packed
    # weight tensor. Alignment will already have been done in the above function.
    weight_tensor = utils.align_width_and_update_state_dict(
        ep, match.weight_node, weight_tensor, align_to=1, force_update=True
    )

    # Also transpose the weight scales tensor to shape [num_groups, N]
    weight_scales_tensor = weight_scales_tensor.transpose(0, 1).contiguous()
    utils.align_width_and_update_state_dict(
        ep,
        match.weight_scales_node,
        weight_scales_tensor,
        align_to=1,
        force_update=True,
    )

    first_graph_node = list(graph_module.graph.nodes)[0]
    with graph_module.graph.inserting_before(first_graph_node):
        weight_tensor_name = utils.get_tensor_name(ep, match.weight_node)
        # Pre-compute the weight sums which are needed to apply activation zero point
        # when using integer accumulation.
        sums_name = weight_tensor_name + "_sums"
        # Sanitize the name
        sums_name = sums_name.replace(".", "_")

        weight_sums_node = create_constant_placeholder(
            exp_program=ep,
            graph=graph_module.graph,
            kind=InputKind.PARAMETER,
            name=sums_name,
            data=sum_per_quant_group,
        )

    with graph_module.graph.inserting_before(match.output_node):
        qlinear_node = graph_module.graph.create_node(
            "call_function",
            exir_ops.edge.et_vk.linear_dq8ca_q4gsw.default,
            args=(
                match.pattern_input_node,
                match.input_scales_node,
                match.input_zeros_node,
                match.weight_node,
                weight_sums_node,
                match.weight_scales_node,
                group_size,
            ),
        )

    qlinear_node.meta["val"] = match.output_node.meta["val"]
    match.output_node.replace_all_uses_with(qlinear_node)


def make_linear_q8ta_q8csw_custom_op(
    ep: ExportedProgram,
    graph_module: torch.fx.GraphModule,
    match: QuantizedLinearMatch,
    weight_tensor: torch.Tensor,
):
    # Pad weight_scales to multiple of 4 so GPU shader reads don't go OOB
    weight_scales_tensor = get_param_tensor(ep, match.weight_scales_node)
    assert weight_scales_tensor is not None
    utils.align_width_and_update_state_dict(
        ep, match.weight_scales_node, weight_scales_tensor
    )

    # Pad bias to multiple of 4 if present
    if match.bias_node is not None:
        bias_tensor = get_param_tensor(ep, match.bias_node)
        if bias_tensor is not None:
            utils.align_width_and_update_state_dict(ep, match.bias_node, bias_tensor)

    first_graph_node = list(graph_module.graph.nodes)[0]
    with graph_module.graph.inserting_before(first_graph_node):
        weight_tensor_name = utils.get_tensor_name(ep, match.weight_node)
        # Pre-compute the weight sums which are needed to apply activation zero point
        # when using integer accumulation.
        sum_per_output_channel = weight_tensor.sum(dim=1).to(torch.int32).contiguous()

        # Pad weight sums to align OC to multiple of 4
        oc = sum_per_output_channel.shape[0]
        if oc % 4 != 0:
            num_padding = 4 - (oc % 4)
            sum_per_output_channel = F.pad(
                sum_per_output_channel, (0, num_padding)
            ).contiguous()

        sums_name = weight_tensor_name + "_sums"
        # Sanitize the name
        sums_name = sums_name.replace(".", "_")

        weight_sums_node = create_constant_placeholder(
            exp_program=ep,
            graph=graph_module.graph,
            kind=InputKind.PARAMETER,
            name=sums_name,
            data=sum_per_output_channel,
        )

    with graph_module.graph.inserting_before(match.output_node):
        qlinear_node = graph_module.graph.create_node(
            "call_function",
            exir_ops.edge.et_vk.linear_q8ta_q8csw.default,
            args=(
                match.pattern_input_node,
                match.input_scales_node,
                match.input_zeros_node,
                match.weight_node,
                weight_sums_node,
                match.weight_scales_node,
            ),
        )

    qlinear_node.meta["val"] = match.output_node.meta["val"]
    match.output_node.replace_all_uses_with(qlinear_node)


def make_q8ta_linear_custom_op(
    ep: ExportedProgram,
    graph_module: torch.fx.GraphModule,
    match: QuantizedLinearMatch,
    weight_tensor: torch.Tensor,
):
    # Pad weight_scales to multiple of 4 so GPU shader reads don't go OOB
    weight_scales_tensor = get_param_tensor(ep, match.weight_scales_node)
    assert weight_scales_tensor is not None
    utils.align_width_and_update_state_dict(
        ep, match.weight_scales_node, weight_scales_tensor
    )

    # Pad bias to multiple of 4 if present
    if match.bias_node is not None:
        bias_tensor = get_param_tensor(ep, match.bias_node)
        if bias_tensor is not None:
            utils.align_width_and_update_state_dict(ep, match.bias_node, bias_tensor)

    first_graph_node = list(graph_module.graph.nodes)[0]
    with graph_module.graph.inserting_before(first_graph_node):
        weight_tensor_name = utils.get_tensor_name(ep, match.weight_node)
        sum_per_output_channel = weight_tensor.sum(dim=1).to(torch.int32).contiguous()

        # Pad weight sums to align OC to multiple of 4
        oc = sum_per_output_channel.shape[0]
        if oc % 4 != 0:
            num_padding = 4 - (oc % 4)
            sum_per_output_channel = F.pad(
                sum_per_output_channel, (0, num_padding)
            ).contiguous()

        sums_name = weight_tensor_name + "_sums"
        sums_name = sums_name.replace(".", "_")

        weight_sums_node = create_constant_placeholder(
            exp_program=ep,
            graph=graph_module.graph,
            kind=InputKind.PARAMETER,
            name=sums_name,
            data=sum_per_output_channel,
        )

    # Use gemv variant when batch size is 1
    input_shape = match.pattern_input_node.meta["val"].shape
    batch_size = input_shape[-2] if len(input_shape) >= 2 else 1
    if batch_size == 1:
        op_target = exir_ops.edge.et_vk.q8ta_linear_gemv.default
    else:
        op_target = exir_ops.edge.et_vk.q8ta_linear.default

    with graph_module.graph.inserting_before(match.output_node):
        qlinear_node = graph_module.graph.create_node(
            "call_function",
            op_target,
            args=(
                match.pattern_input_node,
                match.input_scales_node,
                match.input_zeros_node,
                match.weight_node,
                weight_sums_node,
                match.weight_scales_node,
                match.output_scales_node,
                match.output_zeros_node,
                match.bias_node,
                "relu" if match.relu_node is not None else "none",
            ),
        )

    qlinear_node.meta["val"] = match.quantize_output_node.meta["val"]
    match.quantize_output_node.replace_all_uses_with(qlinear_node)


@register_pattern_replacement("quantized_linear")
def replace_quantized_linear_patterns(
    ep: ExportedProgram,
    graph_module: torch.fx.GraphModule,
    match: QuantizedLinearMatch,
):
    # Extract relevant tensors
    weight_tensor = get_param_tensor(ep, match.weight_node)
    assert weight_tensor is not None

    assert match.weight_scales_node is not None
    weight_scales_tensor = get_param_tensor(ep, match.weight_scales_node)
    assert weight_scales_tensor is not None

    assert match.weight_zeros_node is not None
    weight_zeros_tensor = get_param_tensor(ep, match.weight_zeros_node)
    assert weight_zeros_tensor is not None

    # Route to appropriate custom op.
    # q8ta_linear supports bias, so check it first before the bias guard.
    if (
        match.is_input_static_per_tensor_quantized()
        and match.is_weight_perchannel_quantized()
        and match.has_output_quantization()
    ):
        make_q8ta_linear_custom_op(ep, graph_module, match, weight_tensor)
        return

    # Remaining ops do not support bias
    if match.bias_node is not None:
        return

    if (
        match.is_weight_only_quantized()
        and match.is_weight_pergroup_quantized()
        and utils.is_in_4bit_range(weight_tensor)
    ):
        make_linear_q4gsw_op(
            ep, graph_module, match, weight_tensor, weight_scales_tensor
        )
    elif (
        match.is_input_dynamic_perchannel_quantized()
        and match.is_weight_pergroup_quantized()
        and utils.is_in_4bit_range(weight_tensor)
    ):
        make_linear_dq8ca_q4gsw_op(
            ep, graph_module, match, weight_tensor, weight_scales_tensor
        )
    elif (
        match.is_input_static_per_tensor_quantized()
        and match.is_weight_perchannel_quantized()
    ):
        make_linear_q8ta_q8csw_custom_op(ep, graph_module, match, weight_tensor)