torch/fx/passes/graph_manipulation.py

from typing import Any, Dict, List, NamedTuple, Optional, Tuple

import torch
from torch.fx._compatibility import compatibility
from torch.fx.graph import Graph
from torch.fx.graph_module import GraphModule
from torch.fx.node import (
    _get_qualified_name,
    Argument,
    map_aggregate,
    map_arg,
    Node,
    Target,
)
from torch.fx.passes.param_fetch import lift_lowering_attrs_to_nodes
from torch.fx.passes.shape_prop import ShapeProp


@compatibility(is_backward_compatible=False)
def replace_target_nodes_with(
    fx_module: GraphModule,
    old_op: str,
    old_target: Target,
    new_op: str,
    new_target: Target,
):
    """Modifies all nodes in fx_module.graph.nodes which match the specified op code and target,
    and updates them to match the new op code and target"""
    new_graph = Graph()
    val_map: Dict[Node, Node] = {}
    for node in fx_module.graph.nodes:
        if node.op == old_op and node.target == old_target:
            args = map_arg(node.args, lambda n: val_map[n])
            kwargs = map_arg(node.kwargs, lambda n: val_map[n])
            assert isinstance(args, tuple)
            assert isinstance(kwargs, dict)
            val_map[node] = new_graph.create_node(
                new_op, new_target, args, kwargs, node.name
            )
        else:
            val_map[node] = new_graph.node_copy(node, lambda n: val_map[n])
    fx_module.graph = new_graph


@compatibility(is_backward_compatible=False)
class size_bytes(NamedTuple):
    output_size: int
    total_size: int


@compatibility(is_backward_compatible=False)
def get_size_of_all_nodes(
    fx_module: GraphModule, args: Optional[List[torch.Tensor]] = None
) -> None:
    """Given a fx graph module, update each node with its total size (weights + bias + output)
    and its output_size(output). For a non-module node, the total size is the output size.
    return total size"""
    if args is not None:
        # Mark shape and dtype for each node (node.shape and node.dtype)
        ShapeProp(fx_module).propagate(*args)
    # Calculate the total size of the whole fx graph
    total_size_of_graph = 0.0
    for node in fx_module.graph.nodes:
        if node.op == "output":
            break
        node.size_bytes = get_size_of_node(fx_module, node)
    return


@compatibility(is_backward_compatible=False)
def get_tensor_meta(node: Node) -> Any:
    tensor_meta = node.meta.get("tensor_meta")

    if not tensor_meta:
        raise RuntimeError(
            f"Node {node} has no tensor metadata associated with it! "
            f"Check that shape propagation has run."
        )

    return tensor_meta


@compatibility(is_backward_compatible=False)
def get_size_of_node(fx_module: GraphModule, node: Node) -> size_bytes:
    """Given a node with node.dtype and node.shape, return its total size and its output size.
    total_size = weights + bias + output_size
    """
    # Total num of elements
    total_num_of_elems = 0
    # For a module, conside all parameters
    if node.op == "call_module":
        submodule_dict = dict(fx_module.named_modules())
        submodule = submodule_dict[node.target]
        parameters = submodule.named_parameters()
        # Parameters are named tuples
        for name, p in parameters:
            total_num_of_elems += p.numel()
    # Don't forget the output size
    # node.shape is the shape of this node's output
    tensor_meta = get_tensor_meta(node)
    output_elem = tensor_meta.shape.numel()
    total_num_of_elems += output_elem
    # Assume for now if it's quantized then it's qint8 or quint8
    if tensor_meta.is_quantized:
        size_per_elem_bytes = torch._empty_affine_quantized(
            [], dtype=tensor_meta.dtype
        ).element_size()
    else:
        size_per_elem_bytes = torch.tensor([], dtype=tensor_meta.dtype).element_size()
    total_size = size_per_elem_bytes * total_num_of_elems
    output_size = size_per_elem_bytes * output_elem
    return size_bytes(output_size, total_size)


@compatibility(is_backward_compatible=False)
def serialize_shape(shape: torch.Size) -> str:
    return str(list(shape))


@compatibility(is_backward_compatible=False)
def serialize_stride(stride: Tuple[int]) -> str:
    return str(list(stride))


@compatibility(is_backward_compatible=False)
def serialize_tensor_quantization(
    tensor: torch.Tensor, weights: Dict, pcq_prefix: str
) -> Tuple[Dict, Dict]:
    """
    Args:
        tensor: The tensor from which we try to extract quantization information.
        weights: A dict that contains mapping from name to a tensor value.
        pcq_prefix: A string that we would use later on as prefix for per channel quantization information. This
            usually would be the key that we use to store info of `tensor`.

    Returns:
        scheme: Dict that stores the quantization information of `tensor`.
        per_channel_dict: Dict that stores the information of per_channel_scales and
            per_channel_zero_points of `tensor`. This Will be empty if `tensor` is not
            per channel quantized.

    `tensor` is per tensor quantized:
        scheme: {
            "qscheme": str(tensor.qscheme()),
            "q_scale": tensor.q_scale(),
            "q_zero_point": tensor.q_zero_point(),
        }

    `tensor` is per channel quantized:
        scheme: {
            "qscheme": str(tensor.qscheme()),
            "q_per_channel_scales": {pcq_prefix}_per_channel_scales,
            "q_per_channel_zero_points": {pcq_prefix}_per_channel_zero_points,
            "q_per_channel_axis": tensor.q_per_channel_axis()
        }
        per_channel_dict: {
            {pcq_prefix}_per_channel_scales: {
                "dtype": dtype,
                "shape": shape,
                "is_quantized": is_quantized,
                "stride": stride,
            }
            {pcq_prefix}_per_channel_zero_points: {
                "dtype": dtype,
                "shape": shape,
                "is_quantized": is_quantized,
                "stride": stride,
            }
        }
        weights would be updated with {
            {pcq_prefix}_per_channel_scales: tensor.q_per_channel_scales().float()
            {pcq_prefix}_per_channel_zero_points: tensor.q_per_channel_zero_points().int()
        }
    """
    scheme: Dict[str, Any] = {}
    per_channel_dict: Dict[str, Dict] = {}

    if not tensor.is_quantized:
        return scheme, per_channel_dict

    scheme["qscheme"] = str(tensor.qscheme())

    # For per tensor scheme, we stores scale and zero_point.
    if tensor.qscheme() in {torch.per_tensor_affine, torch.per_tensor_symmetric}:
        scheme["q_scale"] = tensor.q_scale()
        scheme["q_zero_point"] = tensor.q_zero_point()

    # For per channel scheme, per_channel_scales and per_channel_zero_points are tensors.
    # We store their tensor value into `weights` and store the name into `scheme`.
    if tensor.qscheme() in {
        torch.per_channel_affine,
        torch.per_channel_affine_float_qparams,
        torch.per_channel_symmetric,
    }:
        # per_channel_scales is float64. Here we save it as float32.
        weights[
            f"{pcq_prefix}_per_channel_scales"
        ] = tensor.q_per_channel_scales().float()
        scheme["q_per_channel_scales"] = f"{pcq_prefix}_per_channel_scales"
        per_channel_dict.update(
            serialize_weight(
                weights[f"{pcq_prefix}_per_channel_scales"],
                weights,
                f"{pcq_prefix}_per_channel_scales",
            )
        )

        # per_channel_zero_point is int64. Here we save it as int32.
        weights[
            f"{pcq_prefix}_per_channel_zero_points"
        ] = tensor.q_per_channel_zero_points().int()
        scheme["q_per_channel_zero_points"] = f"{pcq_prefix}_per_channel_zero_points"
        per_channel_dict.update(
            serialize_weight(
                weights[f"{pcq_prefix}_per_channel_zero_points"],
                weights,
                f"{pcq_prefix}_per_channel_zero_points",
            )
        )

        scheme["q_per_channel_axis"] = tensor.q_per_channel_axis()
    return scheme, per_channel_dict


@compatibility(is_backward_compatible=False)
def serialize_weight(tensor: torch.Tensor, weights: Dict, name: str) -> Dict:
    weight_dict: Dict[str, Dict] = {name: {}}
    weight_dict[name]["dtype"] = str(tensor.dtype)
    weight_dict[name]["shape"] = serialize_shape(tensor.shape)
    weight_dict[name]["requires_grad"] = str(tensor.requires_grad)
    weight_dict[name]["is_quantized"] = tensor.is_quantized
    weight_dict[name]["stride"] = serialize_stride(tensor.stride())

    if tensor.is_quantized:
        quantization_info, per_channel_dict = serialize_tensor_quantization(
            tensor, weights, name
        )
        weight_dict[name].update(quantization_info)
        weight_dict.update(per_channel_dict)

    return weight_dict


@compatibility(is_backward_compatible=False)
def serialize_leaf_module(
    node: Node, weights_metadata: Dict, weights: Dict, name_prefix: str
) -> Dict:
    parameters: Dict[str, Any] = {}

    for p_name, p_value in node.attrs_for_lowering.items():  # type: ignore[attr-defined]
        if isinstance(p_value, torch.Tensor):
            weights_metadata.update(
                serialize_weight(p_value, weights, f"{name_prefix}.{p_name}")
            )
            weights[f"{name_prefix}.{p_name}"] = p_value
        else:
            parameters[p_name] = str(p_value)

    return parameters


def _update_weight_fused_dtypes(weight, name, node):
    """
    For quantized embedding tables we need to update the shape/type, so we check if the
    users of this get_attr node is a quantized EB and this is the weight for the EB, and
    update the dtype accordingly.
    """
    if len(node.users) == 0:
        return
    user = list(node.users)[0]
    if user.op != "call_function":
        return
    user_target = _get_qualified_name(user.target)
    if (
        user_target.endswith("acc_ops.embedding_bag_byte_rowwise_offsets")
        and node == user.kwargs["weight"]
    ):
        weight[name]["dtype"] = "acc.uint8fused"
    elif (
        user_target.endswith("acc_ops.embedding_bag_4bit_rowwise_offsets")
        and node == user.kwargs["weight"]
    ):
        weight[name]["dtype"] = "acc.uint4fused"


@compatibility(is_backward_compatible=False)
def serialize_module(fx_module: GraphModule, weights: Dict, name_prefix="") -> Dict:
    """Recursively Serializes a graph module (fx_module) to a dictionary which is later exported to JSON.
    It also adds all weights the provided weights dictionary by qualified_name.
    Dictionary Schema:
    MODULE
    {
        modules: {module_name: MODULE],
        nodes: [NODE],
        weights {qualified_name: WEIGHT},
    }
    NODE
    {
        shape: [],
        stride: [],
        dtype: dtype,
        is_quantized: bool,
        target: target,
        op_code: op_code,
        name: name,
        args: [],
        kwargs: {}
    }
    WEIGHT
    {
        dtype: dtype,
        is_quantized: bool,
        shape: [],
        QUANTIZATION,
    }
    QUANTIZATION
    {
        qscheme: qscheme,
        q_scale: float,
        q_zero_point: float,
        q_per_channel_scales, [],
        q_per_channel_zero_points: [],
        q_per_channel_axis, int
    }
    """
    serialized_dict: Dict[str, Any] = {}
    serialized_dict["modules"] = {}
    serialized_dict["weights"] = {}
    serialized_dict["nodes"] = []
    submodules = dict(fx_module.named_modules())
    prefix = f"{name_prefix}." if name_prefix else ""

    def get_node_info(node):
        tensor_meta = get_tensor_meta(node)
        node_rep = {
            "shape": serialize_shape(tensor_meta.shape),
            "dtype": str(tensor_meta.dtype),
            "requires_grad": str(tensor_meta.requires_grad),
            "stride": serialize_stride(tensor_meta.stride),
            "is_quantized": tensor_meta.is_quantized,
        }

        if tensor_meta.is_quantized:
            node_rep["qscheme"] = str(tensor_meta.qparams["qscheme"])

            if tensor_meta.qparams["qscheme"] in {
                torch.per_tensor_affine,
                torch.per_tensor_symmetric,
            }:
                node_rep["q_scale"] = tensor_meta.qparams["scale"]
                node_rep["q_zero_point"] = tensor_meta.qparams["zero_point"]

        # Add all extra lowering_info that was provided in node.meta.
        lowering_info = node.meta.get("lowering_info")
        if lowering_info is not None:
            overlapping_keys = node_rep.keys() & lowering_info.keys()
            assert (
                len(overlapping_keys) == 0
            ), f"Overlap found between lowering_info and node_rep: {overlapping_keys}"
            node_rep.update(lowering_info)

        return node_rep

    # Note: lift_lowering_attrs_to_nodes is only used to support leaf modules
    # that cannot currently be symbolically traced into, e.g. batch norm.
    lift_lowering_attrs_to_nodes(fx_module)
    for node in fx_module.graph.nodes:
        node_rep: Dict[str, Any] = {}
        # Get shape/type info, currently not needed for call_module node
        # whose target is a GraphModule and output node.
        if (
            not (
                node.op == "call_module"
                and isinstance(submodules[node.target], GraphModule)
            )
            and node.op != "output"
        ):
            node_rep.update(get_node_info(node))

        # Recurse down into any submodules we are calling.
        if node.op == "call_module":
            if isinstance(submodules[node.target], GraphModule):
                serialized_module = serialize_module(
                    getattr(fx_module, node.target), weights, node.target
                )
                serialized_dict["modules"][node.target] = serialized_module
            else:
                node_rep["parameters"] = serialize_leaf_module(
                    node,
                    serialized_dict["weights"],
                    weights,
                    prefix + node.target,
                )

        if node.op == "call_function":
            node_rep["target"] = _get_qualified_name(node.target)
        else:
            node_rep["target"] = str(node.target)

        # Make sure we capture all constants.
        if node.op == "get_attr":
            # If we are targeting a parent constant we update the target.
            if node.target.startswith("parent."):
                qualname = node.target[len("parent.") :]
                node.name = qualname
                node_rep["target"] = qualname
            else:
                qualname = prefix + node.target
            # Find the actual target parameter/buffer from the fx_module.
            submod_path, _, target_name = node.target.rpartition(".")
            submod: Optional[torch.nn.Module] = (
                fx_module.get_submodule(submod_path) if submod_path else fx_module
            )
            assert submod is not None, f"submod {submod_path} not found"
            target = getattr(submod, target_name, None)
            assert target is not None, f"{target_name} not an attr of {submod_path}"
            # Check that the target is a tensor, and that we haven't added it already from a leaf module.
            if isinstance(target, torch.Tensor) and qualname not in weights:
                weight = serialize_weight(target, weights, qualname)
                _update_weight_fused_dtypes(weight, qualname, node)
                serialized_dict["weights"].update(weight)
                weights[qualname] = target
        elif node.op == "placeholder":
            ph_type = node.meta.get("ph_type", "")
            assert (
                ph_type == "" or ph_type == "input_ph" or ph_type == "output_ph"
            ), "When present, placeholder type must be 'input_ph' or 'ouput_ph'"
            if ph_type == "input_ph":
                node_rep["ph_type"] = "input_ph"
            elif ph_type == "output_ph":
                node_rep["ph_type"] = "output_ph"

        node_rep["op_code"] = node.op
        node_rep["name"] = node.name

        def get_user_info(user_node: Argument) -> Any:
            return {"is_node": True, "name": str(user_node)}

        def get_arg_info(arg: Argument) -> Any:
            if isinstance(arg, torch.fx.Node):
                return {"is_node": True, "name": str(arg)}
            elif isinstance(arg, (torch.dtype, torch.memory_format, torch.qscheme)):
                return str(arg)
            else:
                return arg

        def get_output_arg_info(arg: Node) -> Dict[str, Any]:
            node_rep: Dict[str, Any] = get_arg_info(arg)
            node_rep.update(get_node_info(arg))
            return node_rep

        if node.op == "output":
            node_rep["args"] = map_arg(
                node.args,
                get_output_arg_info,
            )

            # If there're multiple outputs then node_rep["args"][0] will be a tuple or
            # list. In this case we want to unpack the tuple or list.
            if isinstance(node_rep["args"][0], (tuple, list)):
                node_rep["args"] = node_rep["args"][0]
        else:
            node_rep["args"] = map_aggregate(node.args, get_arg_info)

        node_rep["kwargs"] = map_aggregate(node.kwargs, get_arg_info)
        node_rep["users"] = map_aggregate(list(node.users.keys()), get_user_info)
        serialized_dict["nodes"] += [node_rep]

    return serialized_dict