Repository: pytorch
Issue: Optimizing Per-Channel Quantization for Improved Inference Performance
Labels: oncall: quantization

### 🚀 The feature, motivation and pitch

I have recently been exploring the `torch.export`-based quantization and encountered significant slow-downs in inference performance, particularly with per-channel quantization. To address this, I developed an alternative implementation that leverages broadcasting to enhance inference speed. It achieves about 100X acceleration in inference speed when processing a single image with `resnet50`.

I am enthusiastic about the possibility of contributing this enhancement to the PyTorch community, should it be deemed a valuable addition. Understanding that the team may already be working on similar or alternative solutions, I am open to any feedback or insights you may offer. Your guidance on this matter would be highly appreciated. :)

```python
import time
import torch
import torchvision
import torch.nn.functional as F
from torch.ao.quantization.quantize_pt2e import convert_pt2e, prepare_pt2e
from torch.ao.quantization.quantizer.xnnpack_quantizer import (
    XNNPACKQuantizer,
    get_symmetric_quantization_config,
)


def loop(model, input_):
    with torch.no_grad():
        for i in range(10):
            start = time.time()
            output = model(input_)
            print(f'{time.time() - start:.3f}s')
    return output


def optimized_quantize_per_channel(
        input: torch.Tensor,
        scales: torch.Tensor,
        zero_points: torch.Tensor,
        axis: int,
        quant_min: int,
        quant_max: int,
        dtype: torch.dtype
) -> torch.Tensor:

    shape = [1] * input.dim()
    shape[axis] = -1
    scales = scales.reshape(shape)
    zero_points = zero_points.reshape(shape)

    return torch.clamp(
        torch.round(input / scales) + zero_points,
        quant_min,
        quant_max
    ).to(dtype)


def optimized_dequantize_per_channel(
        input: torch.Tensor,
        scales: torch.Tensor,
        zero_points: torch.Tensor,
        axis: int,
        quant_min: int,
        quant_max: int,
        dtype: torch.dtype
) -> torch.Tensor:

    shape = [1] * input.dim()
    shape[axis] = -1
    scales = scales.reshape(shape)
    zero_points = zero_points.reshape(shape)

    return (input.to(torch.float32) - zero_points) * scales


if __name__ == '__main__':
    inputs = torch.randn(1, 3, 224, 224).cuda()
    model = torchvision.models.resnet50().eval().cuda()

    model = torch._export.capture_pre_autograd_graph(model, [inputs])
    quantizer = XNNPACKQuantizer().set_global(
        get_symmetric_quantization_config(is_per_channel=True)
    )
    model = prepare_pt2e(model, quantizer)
    with torch.no_grad():
        model(inputs)
    model = convert_pt2e(model)


    print('Before optimization')
    output_original = loop(model, inputs)

    # Mock
    torch.ops.quantized_decomposed.quantize_per_channel.default = optimized_quantize_per_channel
    torch.ops.quantized_decomposed.dequantize_per_channel.default = optimized_dequantize_per_channel

    print('\nAfter optimization')
    output_optimized = loop(model, inputs)

    print(f'Cosine similarity: {F.cosine_similarity(output_original.view(1, -1), output_optimized.view(1, -1)).item()}')
```
Sample output from [this Colab](https://colab.research.google.com/drive/1Y7tcqTCajlLqIaZisqqQCrhvJ2vf9Fza?usp=sharing)
```
Before optimization
4.521s
4.337s
3.734s
4.538s
3.784s
3.799s
4.457s
3.802s
3.721s
4.130s

After optimization
0.044s
0.043s
0.040s
0.041s
0.043s
0.040s
0.042s
0.047s
0.039s
0.049s
Cosine similarity: 0.99988853931427
```

Related code:

https://github.com/pytorch/pytorch/blob/0159e3abbd3648498b345af3b6b77e5953ab23b4/torch/ao/quantization/fx/_decomposed.py#L348
https://github.com/pytorch/pytorch/blob/0159e3abbd3648498b345af3b6b77e5953ab23b4/torch/ao/quantization/fx/_decomposed.py#L417


### Alternatives

_No response_

### Additional context

_No response_

cc @jerryzh168 @jianyuh @raghuramank100 @jamesr66a @vkuzo @jgong5 @Xia-Weiwen @leslie-fang-intel