quantize_io_pass.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 logging
from typing import Any, Dict, List, Optional, Union

import numpy as np

import torch

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

from executorch.exir.pass_base import ExportPass
from executorch.exir.tensor import scalar_type_enum
from torch.fx.passes.infra.pass_base import PassResult

logger = logging.getLogger(__name__)


def quantize_input(
    exported_program, input_index, qparams: Optional[Dict[str, Any]] = None
):
    """
    Modify the program to expect quantized input at given index. The input is expected
    to be quantizing this input as the first step. Must be called before
    permute_input_layout. Returns the scale, zero point, qmin, qmax, and dtype of the
    expected quantization.
    """
    graph = exported_program.graph_module.graph
    name = exported_program.graph_signature.user_inputs[input_index]
    placeholders = [n for n in graph.nodes if n.op == "placeholder" and n.name == name]
    assert placeholders
    target_placeholder = placeholders[0]

    if len(target_placeholder.users) != 1:
        raise ValueError(f"Input {input_index} has more than one users")
    quantize = next(iter(target_placeholder.users))
    if (
        quantize.target
        != exir_ops.edge.quantized_decomposed.quantize_per_tensor.default
    ):
        raise ValueError(f"Input {input_index} is not used by a quantize op")

    # If user specified qparams are different from args of quantize op, we do requantization instead of eliminating quantize op
    need_requant = False
    if qparams is not None:
        assert all(
            qparam in qparams for qparam in ["scale", "zp", "dtype"]
        ), "dtype/scale/zp must be specified in qparam for input requantization"
        if qparams["dtype"] != quantize.args[5]:
            if any(
                dtype
                not in [torch.int8, torch.uint8, torch.bool, torch.int16, torch.uint16]
                for dtype in [qparams["dtype"], quantize.args[5]]
            ):
                raise ValueError(
                    f"Only limited data types are supported for requantization, but got {qparams['dtype']} -> {quantize.args[5]}"
                )

            need_requant = True
        elif (
            not np.isclose(qparams["scale"], quantize.args[1])
            or qparams["zp"] != quantize.args[2]
        ):
            need_requant = True

    if need_requant:
        assert qparams is not None
        dtype = qparams["dtype"]
        qmin = torch.iinfo(dtype).min
        qmax = torch.iinfo(dtype).max
        scale = qparams["scale"]
        zero_point = qparams["zp"]
        quant_args = (scale, zero_point, qmin, qmax, dtype)
        logger.info(
            f"Modifying program to requantize quantized input at index {input_index}"
        )
        logger.info(f"Quantization parameters: {quant_args}")

        with exported_program.graph_module.graph.inserting_before(quantize):
            input_dequant = exported_program.graph_module.graph.call_function(
                exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
                args=(
                    target_placeholder,
                    *quant_args,
                ),
            )
            input_dequant.meta["input_qparams"] = [
                {
                    "scale": scale,
                    "zero_point": zero_point,
                    "qmin": qmin,
                    "qmax": qmax,
                    "dtype": dtype,
                }
            ]
            input_dequant.meta["val"] = quantize.meta["val"].to(torch.float32)
            target_placeholder.meta["val"] = target_placeholder.meta["val"].to(dtype)
            quantize.replace_input_with(target_placeholder, input_dequant)
    else:
        quant_args = quantize.args[1:]
        logger.info(f"Modifying program to take quantized input at index {input_index}")
        logger.info(f"Quantization parameters: {quant_args}")

        target_placeholder.meta["val"] = (
            exir_ops.edge.quantized_decomposed.quantize_per_tensor.default(
                target_placeholder.meta["val"], *quant_args
            )
        )
        quantize.replace_all_uses_with(quantize.args[0])

    exported_program.graph_module.graph.eliminate_dead_code()
    return quant_args


def quantize_output(exported_program, output_index):
    """
    Modify the program to produce quantized output at given index. The model is expected
    to be dequantizing this output as the last step. Must be called before
    permute_output_layout. Returns the scale, zero point, qmin, qmax, and dtype of the
    output quantization.
    """
    graph = exported_program.graph_module.graph
    outputs = [n for n in graph.nodes if n.op == "output"]
    if len(outputs) != 1:
        raise NotImplementedError("Only 1 output node is supported")

    output_node = outputs[0]
    output_list = list(output_node.args[0])
    if output_index >= len(output_list):
        raise ValueError(
            f"{len(output_list)} outputs available, "
            + f"output index out of bounds: {output_index}"
        )

    target_output = output_list[output_index]
    if (
        target_output.target
        != exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default
    ):
        raise ValueError("Output {output_index} is not a dequantize op")

    dequant = target_output
    output_list[output_index] = dequant.args[0]
    output_node.args = (output_list,)
    dequant_args = dequant.args[1:]
    graph.eliminate_dead_code()

    logger.info(
        f"Modifying program to produce quantized output at index {output_index}"
    )
    logger.info(f"Dequantization parameters: {dequant_args}")
    return dequant_args


def get_config_method_name(
    prefix: Optional[str] = "forward",
    arg_type: str = "input",
    index: int = 0,
    key: str = "scale",
):
    if prefix is None:
        prefix = ""
    else:
        prefix = prefix + "_"
    assert arg_type in ["input", "output"], "arg_type must be either input or output"
    assert index >= 0, "index must be non-negative"
    assert key in [
        "scale",
        "zp",
        "quant_min",
        "quant_max",
        "dtype",
    ], "key must be one of scale, zp, quant_min, quant_max, dtype"
    return f"{prefix}{arg_type}{index}_{key}"


class QuantizeInputs(ExportPass):
    def __init__(
        self,
        edge_program_manager: EdgeProgramManager,
        quantized_inputs_idx: Union[Dict[int, Dict[str, Any]], List[int]],
        method_name: Optional[str] = None,
    ):
        super().__init__()
        self.edge_program_manager = edge_program_manager

        self.quantized_inputs_idx_dict = {}
        if isinstance(quantized_inputs_idx, dict):
            self.quantized_inputs_idx_dict = quantized_inputs_idx
        else:
            for idx in quantized_inputs_idx:
                self.quantized_inputs_idx_dict[idx] = None
        self.param_prefix_name = method_name

    def call(self, graph_module: torch.fx.GraphModule):
        for i, qparams in self.quantized_inputs_idx_dict.items():
            quant_args = quantize_input(
                self.edge_program_manager.exported_program(), i, qparams
            )

            if not self.edge_program_manager._config_methods:
                self.edge_program_manager._config_methods = {}

            self.edge_program_manager._config_methods[
                get_config_method_name(self.param_prefix_name, "input", i, "scale")
            ] = quant_args[0]
            self.edge_program_manager._config_methods[  # pyre-ignore
                get_config_method_name(self.param_prefix_name, "input", i, "zp")
            ] = quant_args[1]
            self.edge_program_manager._config_methods[
                get_config_method_name(self.param_prefix_name, "input", i, "quant_min")
            ] = quant_args[2]
            self.edge_program_manager._config_methods[
                get_config_method_name(self.param_prefix_name, "input", i, "quant_max")
            ] = quant_args[3]
            self.edge_program_manager._config_methods[
                get_config_method_name(self.param_prefix_name, "input", i, "dtype")
            ] = scalar_type_enum(quant_args[4])
        return PassResult(graph_module, True)


class QuantizeOutputs(ExportPass):
    def __init__(
        self,
        edge_program_manager: EdgeProgramManager,
        quantized_outputs_idx_list: List[int],
        method_name: Optional[str] = None,
    ):
        super().__init__()
        self.edge_program_manager = edge_program_manager
        self.quantized_outputs_idx_list = quantized_outputs_idx_list
        self.param_prefix_name = method_name

    def call(self, graph_module: torch.fx.GraphModule):
        for i in self.quantized_outputs_idx_list:
            dequant_args = quantize_output(
                self.edge_program_manager.exported_program(), i
            )  # noqa F841

            if not self.edge_program_manager._config_methods:
                self.edge_program_manager._config_methods = {}

            self.edge_program_manager._config_methods[
                get_config_method_name(self.param_prefix_name, "output", i, "scale")
            ] = dequant_args[0]
            self.edge_program_manager._config_methods[  # pyre-ignore
                get_config_method_name(self.param_prefix_name, "output", i, "zp")
            ] = dequant_args[1]
            self.edge_program_manager._config_methods[
                get_config_method_name(self.param_prefix_name, "output", i, "quant_min")
            ] = dequant_args[2]
            self.edge_program_manager._config_methods[
                get_config_method_name(self.param_prefix_name, "output", i, "quant_max")
            ] = dequant_args[3]
            self.edge_program_manager._config_methods[
                get_config_method_name(self.param_prefix_name, "output", i, "dtype")
            ] = scalar_type_enum(dequant_args[4])

        return PassResult(graph_module, True)