ai-content-maker/.venv/Lib/site-packages/torch/ao/ns/fx/mappings.py

762 lines
18 KiB
Python

import operator
import torch
import torch.nn as nn
import torch.nn.functional as F
toq = torch.ops.quantized
import torch.ao.nn.quantized as nnq
import torch.ao.nn.quantized.dynamic as nnqd
import torch.ao.nn.intrinsic.quantized as nniq
import torch.ao.nn.intrinsic.quantized.dynamic as nniqd
import torch.ao.nn.intrinsic.qat as nniqat
import torch.ao.nn.intrinsic as nni
import torch.ao.nn.qat as nnqat
import torch.ao.nn.qat.dynamic as nnqatd
from torch.ao.quantization.backend_config import get_native_backend_config
import torch.ao.quantization.fx._lower_to_native_backend as \
_lower_to_native_backend
import torch.ao.quantization.quantization_mappings as quantization_mappings
from .ns_types import NSNodeTargetType
from typing import Callable, Dict, List, Optional, Set, Tuple
def get_base_name_to_sets_of_related_ops() -> Dict[str, Set[NSNodeTargetType]]:
# note: this set is modified below by items from backend_config
sets_of_related_ops: List[Set[NSNodeTargetType]] = [
# conv modules
{
nn.Conv1d,
},
{
nn.Conv2d,
},
{
nn.Conv3d,
},
# conv functionals
{
F.conv1d,
},
{
F.conv2d,
},
{
F.conv3d,
},
# linear modules
{
nn.Linear,
},
# linear functionals
{
F.linear,
},
# average pool
{
nn.AvgPool1d,
torch.avg_pool1d,
},
{
nn.AvgPool2d,
torch._C._nn.avg_pool2d,
},
{
nn.AvgPool3d,
torch._C._nn.avg_pool3d,
},
# adaptive average pool
{
nn.AdaptiveAvgPool1d,
F.adaptive_avg_pool1d,
},
{
nn.AdaptiveAvgPool2d,
F.adaptive_avg_pool2d,
},
{
nn.AdaptiveAvgPool3d,
F.adaptive_avg_pool3d,
},
# LSTM
{
nn.LSTM,
},
# add
{
torch.add,
operator.add, # x + y
},
# cat
{
torch.cat,
},
# mul
{
torch.mul,
operator.mul,
},
# relu
{
F.relu,
nn.ReLU,
'relu',
'relu_',
torch.relu,
},
# maxpool
{
nn.MaxPool1d,
F.max_pool1d,
},
{
nn.MaxPool2d,
F.max_pool2d,
},
{
nn.MaxPool3d,
F.max_pool3d,
},
# sigmoid
{
torch.sigmoid,
'sigmoid',
'sigmoid_',
nn.Sigmoid,
F.sigmoid,
},
# BatchNorm
{
nn.BatchNorm2d,
},
{
nn.BatchNorm3d,
},
# ConvTranspose
{
nn.ConvTranspose1d,
},
{
nn.ConvTranspose2d,
},
{
nn.ConvTranspose3d,
},
# functional transposed conv
{
F.conv_transpose1d,
},
{
F.conv_transpose2d,
},
{
F.conv_transpose3d,
},
# ELU
{
nn.ELU,
},
# Embedding
{
nn.Embedding,
},
# EmbeddingBag
{
nn.EmbeddingBag,
},
# GroupNorm
{
nn.GroupNorm,
},
# Hardswish
{
nn.Hardswish,
},
# InstanceNorm
{
nn.InstanceNorm1d,
},
{
nn.InstanceNorm2d,
},
{
nn.InstanceNorm3d,
},
# LayerNorm
{
nn.LayerNorm,
},
# LeakyReLU
{
nn.LeakyReLU,
},
# ReLU6
{
nn.ReLU6,
F.relu6,
},
# F.elu
{
F.elu,
},
# F.hardswish
{
F.hardswish,
},
# F.group_norm
{
F.group_norm,
},
# F.instance_norm
{
F.instance_norm,
},
# F.layer_norm
{
F.layer_norm,
},
# F.leaky_relu
{
F.leaky_relu,
},
# F.silu
{
nn.SiLU,
F.silu,
},
# F.mish
{
nn.Mish,
F.mish,
},
# F.tanh
{
nn.Tanh,
F.tanh,
torch.tanh,
'tanh_',
'tanh',
},
# F.hardsigmoid
{
'hardsigmoid_',
'hardsigmoid',
F.hardsigmoid,
nn.Hardsigmoid,
},
# F.hardtanh
{
nn.Hardtanh,
F.hardtanh,
F.hardtanh_,
},
# floordiv
{
operator.floordiv,
},
# unsqueeze
{
torch.unsqueeze,
},
# stack
{
torch.stack,
},
# squeeze
{
torch.squeeze,
},
# sort
{
torch.sort,
},
# repeat_interleave
{
torch.repeat_interleave,
},
# min
{
torch.min,
},
# mean
{
torch.mean,
},
# max
{
torch.max,
},
# transpose
{
torch.transpose,
},
# flatten
{
torch.flatten,
},
# clamp
{
torch.clamp,
},
# chunk
{
torch.chunk,
},
# interpolate
{
torch.nn.functional.interpolate,
},
# dropout
{
nn.Dropout,
},
# F.dropout
{
F.dropout,
},
# matmul
{
torch.matmul,
},
# Softmax
{
nn.Softmax,
},
# PReLU
{
nn.PReLU,
nnq.PReLU,
},
# F.prelu
{
F.prelu,
toq.prelu,
},
# pixel shuffle
{
nn.PixelShuffle,
},
{
F.pixel_shuffle,
},
# pixel unshuffle
{
nn.PixelUnshuffle,
},
{
F.pixel_unshuffle,
},
# narrow
{
torch.narrow,
},
]
# for each floating point op, add versions of the op added by
# backend_config
backend_config = get_native_backend_config()
new_connections: List[Tuple[Callable, Callable]] = [
# technical debt edge case
(nn.Linear, nn.modules.linear.NonDynamicallyQuantizableLinear),
]
for pattern, config in backend_config._pattern_complex_format_to_config.items():
# pattern format: (c, (b, a))
first_element = pattern
# look from the end, because pattern is in reverse order
while isinstance(first_element, (list, tuple)):
first_element = first_element[-1]
if config.fused_module is not None:
# case 1: pattern fuses a pattern of ops into an op
# example: nn.Conv1d, nn.ReLU fused into nni.ConvReLU1d
new_connections.append((first_element, config.fused_module))
if config.qat_module is not None:
# case 2: pattern swaps a module into a QAT module
# example: nni.ConvReLU1d swapped into nniqat.ConvReLU1d
new_connections.append((first_element, config.qat_module))
if config.reference_quantized_module is not None:
# case 3: reference version of floating point module, such as
# nn.Conv2d and nnqr.Conv2d
new_connections.append((first_element, config.reference_quantized_module))
#
# Add reference module swaps from default lowering path
#
for source_to_target in (
_lower_to_native_backend.STATIC_LOWER_MODULE_MAP,
_lower_to_native_backend.DYNAMIC_LOWER_MODULE_MAP,
_lower_to_native_backend.WEIGHT_ONLY_LOWER_MODULE_MAP,
_lower_to_native_backend.SPECIAL_PATTERN_LOWER_MODULE_MAP,
):
for source, target in source_to_target.items(): # type: ignore[attr-defined]
new_connections.append((source, target))
for source_to_double_target in (
_lower_to_native_backend.STATIC_LOWER_FUSED_MODULE_MAP,
_lower_to_native_backend.STATIC_LOWER_FUSED_MODULE_TWO_INPUTS_MAP,
_lower_to_native_backend.DYNAMIC_LOWER_FUSED_MODULE_MAP,
):
for source, (target1, target2) in source_to_double_target.items(): # type: ignore[attr-defined]
new_connections.append((source, target1))
new_connections.append((source, target2))
#
# Add function swaps from default lowering path
#
for source, (target1, target2) in \
_lower_to_native_backend.STATIC_LOWER_FUNCTIONAL_MAP.items():
new_connections.append((source, target1))
new_connections.append((source, target2))
for source_to_target in (
_lower_to_native_backend.QBIN_OP_MAPPING,
_lower_to_native_backend.QBIN_RELU_OP_MAPPING,
quantization_mappings.DEFAULT_FLOAT_TO_QUANTIZED_OPERATOR_MAPPINGS,
):
for source, target in source_to_target.items():
new_connections.append((source, target))
#
# Add other swaps, ideally in the future this could be removed
# after the lowering code stops using these.
#
for source_to_target in (
quantization_mappings.DEFAULT_DYNAMIC_QUANT_MODULE_MAPPINGS,
):
for source, target in source_to_target.items():
new_connections.append((source, target))
# add the new connections from backend_config
for item1, item2 in new_connections:
for set_of_related_ops in sets_of_related_ops:
if item1 in set_of_related_ops or item2 in set_of_related_ops:
set_of_related_ops.add(item1)
set_of_related_ops.add(item2)
break
base_name_to_sets_of_related_ops: Dict[str, Set[NSNodeTargetType]] = {}
counter = 0
for set_of_related_ops in sets_of_related_ops:
base_name = str(counter)
counter += 1
base_name_to_sets_of_related_ops[base_name] = set_of_related_ops
return base_name_to_sets_of_related_ops
def get_base_name_for_op(
base_name_to_sets_of_related_ops: Dict[str, Set[NSNodeTargetType]],
op: NSNodeTargetType,
) -> Optional[str]:
for base_name, set_of_related_ops in base_name_to_sets_of_related_ops.items():
if op in set_of_related_ops:
return base_name
return None
def add_op_to_sets_of_related_ops(
base_name_to_sets_of_related_ops: Dict[str, Set[NSNodeTargetType]],
op: NSNodeTargetType,
related_op: Optional[NSNodeTargetType],
) -> None:
if related_op is not None:
for set_of_related_ops in base_name_to_sets_of_related_ops.values():
if related_op in set_of_related_ops:
set_of_related_ops.add(op)
return
# if we got here, related_op was not found
raise AssertionError(f"{related_op} was not found")
else:
counter = 0
while str(counter) in base_name_to_sets_of_related_ops:
counter += 1
base_name_to_sets_of_related_ops[str(counter)] = {op}
# TODO(future PR): clean this up
def get_node_type_to_io_type_map() -> Dict[str, Set[NSNodeTargetType]]:
FUNS_IO_TYPE_FP32: Set[NSNodeTargetType] = {
F.linear,
F.conv1d,
F.conv2d,
F.conv3d,
torch.cat,
F.elu,
F.hardswish,
F.instance_norm,
F.layer_norm,
F.leaky_relu,
F.dropout,
F.silu,
F.mish,
operator.add,
torch.add,
operator.mul,
torch.mul,
torch.sum,
F.prelu,
}
FUNS_IO_TYPE_FP16: Set[NSNodeTargetType] = set()
FUNS_IO_TYPE_INT8: Set[NSNodeTargetType] = {
toq.linear,
toq.linear_relu,
toq.conv1d,
toq.conv1d_relu,
toq.conv2d,
toq.conv2d_relu,
toq.conv3d,
toq.conv3d_relu,
toq.cat,
toq.elu,
toq.hardswish,
toq.instance_norm,
toq.layer_norm,
toq.leaky_relu,
toq.dropout,
toq.prelu,
# TODO(future PR): implement shadowing for binary ops and
# uncomment below
# toq.add,
# toq.mul,
}
FUNS_IO_TYPE_FP32_OR_INT8: Set[NSNodeTargetType] = {
F.relu,
F.tanh,
torch.tanh,
F.sigmoid,
torch.sigmoid,
F.hardsigmoid,
operator.floordiv,
torch.adaptive_avg_pool1d,
F.adaptive_avg_pool2d,
F.adaptive_avg_pool3d,
F.dropout,
F.hardtanh,
F.hardtanh_,
F.interpolate,
F.max_pool1d,
F.max_pool2d,
F.max_pool3d,
F.relu6,
F.pixel_shuffle,
F.pixel_unshuffle,
torch.avg_pool1d,
torch._C._nn.avg_pool2d,
torch._C._nn.avg_pool3d,
torch.cat,
torch.chunk,
torch.clamp,
torch.flatten,
torch.transpose,
torch.max,
torch.mean,
torch.min,
torch.narrow,
torch.repeat_interleave,
torch.sort,
torch.squeeze,
torch.stack,
torch.unsqueeze,
operator.add,
}
MODS_IO_TYPE_FP32: Set[NSNodeTargetType] = {
nn.Linear,
nnqat.Linear,
nnqatd.Linear,
nnqd.Linear,
torch.nn.modules.linear.NonDynamicallyQuantizableLinear,
nn.Conv1d,
nn.Conv2d,
nn.Conv3d,
nnqat.Conv1d,
nnqat.Conv2d,
nnqat.Conv3d,
nnqat.Embedding,
nnqat.EmbeddingBag,
nn.LSTM,
# note: nnqd.Linear is an instance of nnq.Linear, so this
# check has to happen before the int8 module check
nnqd.LSTM,
nn.BatchNorm2d,
nn.BatchNorm3d,
nn.Dropout,
nn.ConvTranspose1d,
nn.ConvTranspose2d,
nn.ConvTranspose3d,
nn.ELU,
nn.GroupNorm,
nn.InstanceNorm1d,
nn.InstanceNorm2d,
nn.InstanceNorm3d,
nn.LayerNorm,
nn.Hardswish,
nn.LeakyReLU,
nn.ReLU6,
nn.SiLU,
nn.Mish,
nn.Softmax,
nn.PReLU,
nni.BNReLU2d,
nni.BNReLU3d,
nni.ConvReLU1d,
nni.ConvReLU2d,
nni.ConvReLU3d,
nni.LinearReLU,
nni.LinearBn1d,
nni.ConvBn1d,
nni.ConvBn2d,
nni.ConvBn3d,
nniqat.ConvBn1d,
nniqat.ConvBn2d,
nniqat.ConvBn3d,
nniqat.ConvBnReLU1d,
nniqat.ConvBnReLU2d,
nniqat.ConvBnReLU3d,
nniqat.ConvReLU1d,
nniqat.ConvReLU2d,
nniqat.ConvReLU3d,
nniqat.LinearReLU,
nniqat.LinearBn1d,
nniqd.LinearReLU,
nni.LinearLeakyReLU,
nni.LinearTanh,
nni.ConvAdd2d,
nni.ConvAddReLU2d,
}
MODS_IO_TYPE_INT8: Set[NSNodeTargetType] = {
nnq.Linear,
nnq.Conv1d,
nnq.Conv2d,
nnq.Conv3d,
nnq.BatchNorm2d,
nnq.BatchNorm3d,
nnq.Dropout,
nnq.ConvTranspose1d,
nnq.ConvTranspose2d,
nnq.ELU,
nnq.InstanceNorm1d,
nnq.InstanceNorm2d,
nnq.InstanceNorm3d,
nnq.LayerNorm,
nnq.Hardswish,
nnq.LeakyReLU,
nnq.Embedding,
nnq.EmbeddingBag,
nnq.Dropout,
nnq.Softmax,
nnq.PReLU,
nniq.BNReLU2d,
nniq.BNReLU3d,
nniq.ConvReLU1d,
nniq.ConvReLU2d,
nniq.ConvReLU3d,
nniq.LinearReLU,
nniq.LinearLeakyReLU,
nniq.LinearTanh,
nniq.ConvAdd2d,
nniq.ConvAddReLU2d,
}
MODS_IO_TYPE_FP32_OR_INT8: Set[NSNodeTargetType] = {
nn.ReLU,
nn.Tanh,
nn.Sigmoid,
nn.Hardsigmoid,
nn.AdaptiveAvgPool1d,
nn.AdaptiveAvgPool2d,
nn.AdaptiveAvgPool3d,
nn.AvgPool1d,
nn.AvgPool2d,
nn.AvgPool3d,
nn.Dropout,
nn.Hardtanh,
nn.Identity,
nn.MaxPool1d,
nn.MaxPool2d,
nn.MaxPool3d,
nn.PixelShuffle,
nn.PixelUnshuffle,
nn.ReLU6,
}
METHS_IO_TYPE_FP32_OR_INT8: Set[NSNodeTargetType] = {
'sigmoid_',
'sigmoid',
'tanh_',
'tanh',
'hardsigmoid_',
'hardsigmoid',
'relu_',
'relu',
}
return {
'funs_io_type_fp32': FUNS_IO_TYPE_FP32,
'funs_io_type_fp16': FUNS_IO_TYPE_FP16,
'funs_io_type_int8': FUNS_IO_TYPE_INT8,
'funs_io_type_fp32_or_int8': FUNS_IO_TYPE_FP32_OR_INT8,
'mods_io_type_fp32': MODS_IO_TYPE_FP32,
'mods_io_type_int8': MODS_IO_TYPE_INT8,
'mods_io_type_fp32_or_int8': MODS_IO_TYPE_FP32_OR_INT8,
'meths_io_type_fp32_or_int8': METHS_IO_TYPE_FP32_OR_INT8,
}
def get_unmatchable_types_map() -> Dict[str, Set[NSNodeTargetType]]:
FUNS_UNMATCHABLE: Set[NSNodeTargetType] = {
torch.quantize_per_tensor,
operator.getitem,
}
MODS_UNMATCHABLE: Set[NSNodeTargetType] = {
nn.Identity,
}
METHS_UNMATCHABLE: Set[NSNodeTargetType] = {
'to',
'dequantize',
'reshape',
'view',
'unsqueeze_',
'unsqueeze',
'transpose',
'squeeze_',
'squeeze',
'size',
'shape',
'resize_',
'repeat_interleave',
'repeat',
'permute',
'numel',
'mean',
'detach_',
'detach',
'contiguous',
'clamp',
'chunk',
}
return {
'funs_unmatchable': FUNS_UNMATCHABLE,
'mods_unmatchable': MODS_UNMATCHABLE,
'meths_unmatchable': METHS_UNMATCHABLE,
}