[RFC] TorchInductor with X86 CPU as backend of Quantization in PyTorch 2.0 Export
Labels: oncall: quantization, triaged, oncall: pt2

## 🚀 The feature, motivation and pitch

This RFC proposes to add [TorchInductor](https://dev-discuss.pytorch.org/t/torchinductor-a-pytorch-native-compiler-with-define-by-run-ir-and-symbolic-shapes/747) with X86 CPU device as one of the backends for [Quantization 2.0 in Export](https://pytorch.org/tutorials/prototype/pt2e_quant_ptq_static.html).

* New Quantization 2.0 flow uses the PT2 Export workflow (`torch._dynamo.export`) to capture the model into a graph and perform quantization transformations on top of the ATen dialect graph. This approach is expected to have significantly higher model coverage, better programmability, and a simplified UX.
* TorchInductor is the new compiler backend that compiles the FX Graphs generated by TorchDynamo into optimized C++/Triton kernels.
* This RFC focus on X86 CPU device to combine Quantization 2.0 flow and TorchInductor. As for the support of other devices, please refer to [open discussion section](#open-discussion)

The proposed high-level architecture of quantization 2.0 with Inductor could look like this:

```
    float_model(Python)                               Input
        \                                              /
         \                                            /
    —-------------------------------------------------------
    |                    Dynamo Export                     |
    —-------------------------------------------------------
                                |
                        FX Graph in ATen     
                                |            X86InductorQuantizer
                                |                 /
                                |                /
    —--------------------------------------------------------
    |                 prepare_pt2e                       |
    —--------------------------------------------------------
                                |
                         Calibrate/Train
                                |
    —--------------------------------------------------------
    |                      convert_pt2e                     |
    —--------------------------------------------------------
                                |
                    Reference Quantized Model
                                |
    —--------------------------------------------------------
    |                        Lowering                       |
    —--------------------------------------------------------
                                |
                            Inductor
```

The proposed UX is as below:
```
import torch
import torch._dynamo as torchdynamo
from torch.ao.quantization._quantize_pt2e import convert_pt2e, prepare_pt2e_quantizer
import torch.ao.quantization._pt2e.quantizer.x86_inductor_quantizer as xiq

class M(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(5, 10)

    def forward(self, x):
        return self.linear(x)

example_inputs = (torch.randn(1, 5),)
model = M().eval()

# Step 1: Trace the model into an FX graph of flattened ATen operators
exported_graph_module, guards = torchdynamo.export(
    model,
    *copy.deepcopy(example_inputs),
    aten_graph=True,
)

# Step 2: Insert observers or fake quantize modules
quantizer = xiq.X86InductorQuantizer()
operator_config = xiq.get_symmetric_quantization_config(is_per_channel=True)
quantizer.set_global(operator_config)
prepared_graph_module = prepare_pt2e(exported_graph_module, quantizer)
# Doing calibration here.

# Step 3: Quantize the model
convert_graph_module = convert_pt2e(prepared_graph_module)

# Step 4: Lower Reference Quantized Model into the backend
compile_model = torch.compile(convert_graph_module)
```

## Alternatives

### Frontend Changes
The frontend will follow the design of Quantizer as in [PyTorch 2.0 Export Post Training Static Quantization](https://pytorch.org/tutorials/prototype/pt2e_quant_ptq_static.html).

* `prepare_pt2e` and `convert_pt2e` are the APIs provided by standard Quantization 2.0 flow.
* `X86InductorQuantizer` will be enabled for the quantization recipes on X86 platform as in this draft [PR-98730](https://github.com/pytorch/pytorch/pull/98730). Although, the quantization annotation API is not finalized, `X86InductorQuantizer` can follow the changes in quantization annotation API.
* `Quantized Model Representation` is the IR user saw after the quantization flow of `convert_pt2e`. It's important to the backend developer for doing the quantization fusion and lowering. The final design and implementation of `Quantized Model Representation` are not ready, we will rely on current `Quantized Model Representation` to enable quantization fusion and kernel code-gen inside Inductor for now. The current `Quantized Model Representation` is as:
    * Take convolution as an example, the quantization pattern is like `dequantize_per_tensor -> fp32_convolution -> quant_per_tensor`
    * `dequantize_per_tensor` and `quantize_per_tensor` are not decomposed into primary aten operators.
We will match the quantization patterns based on the above `Quantized Model Representation` inside inductor for now. Meanwhile, the change of `Quantized Model Representation` will gradually take effect per operator. The backend pattern matcher will follow up with the changes of `Quantized Model Representation` gradually.

### Backend Changes
We will use the `torch.compile` API to lower the `Quantized Model Representation` into inductor. Since the `Quantized Model Representation` is captured by `torch._dynamo.export`, we probably don't need to use the `torch.compile` API which will trigger the `dynamo` process for FX graph capture again.
Inside inductor, we will enable the external-call for computation-intensive operators/patterns (convolution, linear) and C++ backend code-gen for memory-intensive operators/patterns.

#### Computation-intensive operators/patterns
For computation-intensive operators/patterns (convolution, linear), we will match the quantization pattern and lower it into Inductor external-call.

##### Pattern Matcher
Refer to the [PR-101164](https://github.com/pytorch/pytorch/pull/101164) for general implementations of the pattern matcher and external-call lowering process:

* `dequantize_per_tensor` and `quantize_per_tensor` will be decomposed into primary operators like `mul`, `div`, `to` inside AOT-Autograd.
* We will use Inductor `register_lowering_pattern` function to match `decomposed dequantize_per_tensor -> aten.convolution -> decomposed quantize_per_tensor` pattern and substitute it into a `QConv` External Call. `QConv` External Call will codegen the C++ code with lines of new quantized convolution operators implementation.
* As we mentioned in the `Frontend` section of `Quantized Model Representation`, `Quantized Model Representation` is not finalized and the quantization patterns that Inductor saw will be subject to change accordingly.

##### New quantized operators implementation
We plan to enable new quantized operators implementation for quantization 2.0 and put them in PyTorch to support Inductor backend as we did for FP32 and BF16 operators in Inductor.

* Previously, we have [quantized convolution operators](https://github.com/pytorch/pytorch/blob/e3ee5b00beff9401d69206b84b82323f7a63a048/aten/src/ATen/native/quantized/library.cpp#L66) which accepts quantized tensors of activation, weight as inputs. Quantized activation has the scale and zero_point information which is needed for the quantized convolution calculation. For quantization 2.0, we will only see the plain tensor of uint8/int8 data type, so the schema of the new quantized convolution operator is subject to be changed for accepting extra input of scale, zero_point for activation, weight, and output tensors.
* Previously, the weight of convolution is encapsulated in a `Conv2dPackedParamsBase` object instead of a tensor. Also, the qconv's implementation is encapsulated inside the `PackedConvWeight` class. Some runtime information for the calculation can be fetched from the attr of the `PackedConvWeight` object. For quantization 2.0, we will implement the new quantized convolution as functional invoking with only Tensor objects in inputs.

Due to these differences in 1.X and 2.X path, we need to enable new quantized operators implementation.

##### Weight prepack implementation
OneDNN convolution needs prepacked weight to do the calculation. The weight prepack can be done in the graph preparation phase instead of the runtime to avoid runtime overhead. As we discuss [here](https://github.com/pytorch/pytorch/pull/101164#discussion_r1195953417), to do weight prepack we need the constant fold feature support in Inductor.

* Current constant folding feature doesn't support FX graph captured by `torch._dynamo.export` as reported in [issue-103582](https://github.com/pytorch/pytorch/issues/103582). It needs further discussion about how to resolve this issue.
* Current constant folding feature will fold constants when possible. It will fold `fp32_weight->quant_per_channel->dequant_per_channel` into a new `fp32` weight as discussed [here](https://github.com/pytorch/pytorch/pull/100652#discussion_r1218789255). We need to keep `dequant_per_channel` node in the graph. Then `fp32_weight->quant_per_channel` can be constant folding into an `int8_weight`, and `dequant_per_channel->conv` can be further lowered into an `int8_conv` node. 

Suppose we resolve the above issues, the design of weight prepack feature can be:
* Step 1: `fp32_weight->quant_per_channel` can be constant folding into an `int8_weight`.
* Step 2: Doing weight prepack of this `int8_weight` to pack it into an `int8 Mkldnn Tensor`. Meanwhile, the `dequant_per_channel->aten.convolution` pattern will be substituted with a `dynamic_quant_convloution` node which accepts `fp32 tensor activation` and `int8 Mkldnn Tensor weight` as inputs.
* Step 3: Pattern match the pattern of
```

decomposed_dequantize_per_tensor  int8_mkldnn_tensor_weight
    \                                   / 
        dynamic_quant_convloution
                 |
    optional(decomposed_quantize_per_tensor) 
```
and substitue the pattern into the `QConv` external call node.

##### Primitive Cache Design
OneDNN Primitive Cache is a feature for Quantized 1.X FX path of oneDNN convolution implementation. It can eliminate the overhead of re-creation of oneDNN primitive. For 2.X path, it's still a design open to be explored further.

##### CPPWrap Support
CPPWrap can reduce the python overhead for inductor generated code. Quantization also needs the support of this feature. Since CPPWrap already supports fp32 and bf16 external calls, if the quantization follows the designs of fp32 and bf16 external calls. CPPWrap should be able to support it by nature.

#### Memory-intensive operators
As for the Memory-intensive operators besides convolution and linear, we will reply on the inductor C++ backend code-gen capability. Typical, there are 3 patterns will see in quantization.

##### Dequant->op code gen
As we mentioned above, `decomposed dequantize_per_tensor` has been decomposed into primary operators of `to_fp32`, `sub`, and `mul` after this [PR-99131](https://github.com/pytorch/pytorch/pull/99131). These primary operators support code-gen and loop fusions with the following memory-intensive operator.

Take example of `dequantize_per_tensor->relu` pattern as example.
```
def fn(x):
    tmp = torch.ops.quantized_decomposed.dequantize_per_tensor.tensor(
        x,
        scale=torch.tensor(0.1, dtype=torch.float),
        zero_point=torch.tensor(1, dtype=torch.int64),
        quant_min=0,
        quant_max=255,
        dtype=torch.uint8,
    )
    y = torch.relu(tmp)
    return y
```
The generated code on AVX-512 supported CPU will be:
```
cpp_fused_dequantize_per_tensor_lift_fresh_relu_0 = async_compile.cpp('''
#include "/tmp/torchinductor_root/mq/cmqzxwuyo7ryvun3egqos5jq5ak4fue7d2jbopbqs7pgpkhdpfh4.h"
extern "C" void kernel(const unsigned char* in_ptr0,
                       float* out_ptr0)
{
    {
        for(long i0=static_cast<long>(0L); i0<static_cast<long>(198144L); i0+=static_cast<long>(16L))
        {
            auto tmp0 = at::vec::load_uint8_as_float(in_ptr0 + static_cast<long>(i0));
            auto tmp1 = (tmp0);
            auto tmp2 = at::vec::Vectorized<float>(static_cast<float>(1.0));
            auto tmp3 = tmp1 - tmp2;
            auto tmp4 = at::vec::Vectorized<float>(static_cast<float>(0.10000000149011612));
            auto tmp5 = tmp3 * tmp4;
            auto tmp6 = at::vec::clamp_min(tmp5, decltype(tmp5)(0));
            tmp6.store(out_ptr0 + static_cast<long>(i0));
        }
        #pragma omp simd simdlen(8)
        for(long i0=static_cast<long>(198144L); i0<static_cast<long>(198147L); i0+=static_cast<long>(1L))
        {
            auto tmp0 = in_ptr0[static_cast<long>(i0)];
            auto tmp1 = static_cast<float>(tmp0);
            auto tmp2 = static_cast<float>(1.0);
            auto tmp3 = tmp1 - tmp2;
            auto tmp4 = static_cast<float>(0.10000000149011612);
            auto tmp5 = decltype(tmp3)(tmp3 * tmp4);
            auto tmp6 = tmp5 * (tmp5>0);
            out_ptr0[static_cast<long>(i0)] = tmp6;
        }
    }
}
''')

```

##### op->quant code gen
As we mentioned above, `decomposed quantize_per_tensor` has been decomposed into primary operators of `div`, `mul` and `add`, `to_uint8` after this [PR-99131](https://github.com/pytorch/pytorch/pull/99131). These primary operators support code-gen and loop fusions with the following memory-intensive operator.

Take an example of `relu->quantize_per_tensor` pattern:
```
def fn(x):
    tmp = torch.relu(x)
    y = torch.ops.quantized_decomposed.quantize_per_tensor.tensor(
        tmp,
        scale=torch.tensor(0.1, dtype=torch.float),
        zero_point=torch.tensor(1, dtype=torch.int64),
        quant_min=0,
        quant_max=255,
        dtype=torch.uint8,
    )
    return y
```
The generated code on AVX-512 supported CPU will be:
```
cpp_fused_lift_fresh_quantize_per_tensor_relu_0 = async_compile.cpp('''
#include "/tmp/torchinductor_root/mq/cmqzxwuyo7ryvun3egqos5jq5ak4fue7d2jbopbqs7pgpkhdpfh4.h"
extern "C" void kernel(const float* in_ptr0,
                       unsigned char* out_ptr0)
{
    {
        for(long i0=static_cast<long>(0L); i0<static_cast<long>(198144L); i0+=static_cast<long>(16L))
        {
            auto tmp0 = at::vec::Vectorized<float>::loadu(in_ptr0 + static_cast<long>(i0));
            auto tmp1 = at::vec::clamp_min(tmp0, decltype(tmp0)(0));
            auto tmp2 = at::vec::Vectorized<float>(static_cast<float>(0.10000000149011612));
            auto tmp3 = tmp2.reciprocal();
            auto tmp4 = at::vec::Vectorized<float>(static_cast<float>(1.0));
            auto tmp5 = tmp3 * tmp4;
            auto tmp6 = tmp1 * tmp5;
            auto tmp7 = tmp6.round();
            auto tmp8 = tmp7 + tmp4;
            auto tmp9 = at::vec::Vectorized<float>(static_cast<float>(0.0));
            auto tmp10 = at::vec::maximum(tmp8, tmp9);
            auto tmp11 = at::vec::Vectorized<float>(static_cast<float>(255.0));
            auto tmp12 = at::vec::minimum(tmp10, tmp11);
            auto tmp13 = (tmp12);
            at::vec::store_float_as_uint8(tmp13, out_ptr0 + static_cast<long>(i0));
        }
        #pragma omp simd simdlen(8)
        for(long i0=static_cast<long>(198144L); i0<static_cast<long>(198147L); i0+=static_cast<long>(1L))
        {
            auto tmp0 = in_ptr0[static_cast<long>(i0)];
            auto tmp1 = tmp0 * (tmp0>0);
            auto tmp2 = static_cast<float>(0.10000000149011612);
            auto tmp3 = 1 / tmp2;
            auto tmp4 = static_cast<float>(1.0);
            auto tmp5 = decltype(tmp3)(tmp3 * tmp4);
            auto tmp6 = decltype(tmp1)(tmp1 * tmp5);
            auto tmp7 = std::nearbyint(tmp6);
            auto tmp8 = tmp7 + tmp4;
            auto tmp9 = static_cast<float>(0.0);
            auto tmp10 = max_propagate_nan(tmp8, tmp9);
            auto tmp11 = static_cast<float>(255.0);
            auto tmp12 = min_propagate_nan(tmp10, tmp11);
            auto tmp13 = static_cast<unsigned char>(tmp12);
            out_ptr0[static_cast<long>(i0)] = tmp13;
        }
    }
}
''')

```

##### Dequant->op->quant code gen
This case is a combination of the above 2 cases, take an example of `dequantize_per_tensor->relu->quantize_per_tensor`:
```
def fn(x):
    tmp = torch.ops.quantized_decomposed.dequantize_per_tensor.tensor(
        x,
        scale=torch.tensor(0.1, dtype=torch.float),
        zero_point=torch.tensor(1, dtype=torch.int64),
        quant_min=0,
        quant_max=255,
        dtype=torch.uint8,
    )
    tmp = torch.relu(tmp)
    y = torch.ops.quantized_decomposed.quantize_per_tensor.tensor(
        tmp,
        scale=torch.tensor(0.1, dtype=torch.float),
        zero_point=torch.tensor(1, dtype=torch.int64),
        quant_min=0,
        quant_max=255,
        dtype=torch.uint8,
    )
    return y
```
The generated code should be:
```
cpp_fused_dequantize_per_tensor_lift_fresh_quantize_per_tensor_relu_0 = async_compile.cpp('''
#include "/tmp/torchinductor_root/mq/cmqzxwuyo7ryvun3egqos5jq5ak4fue7d2jbopbqs7pgpkhdpfh4.h"
extern "C" void kernel(const unsigned char* in_ptr0,
                       unsigned char* out_ptr0)
{
    {
        for(long i0=static_cast<long>(0L); i0<static_cast<long>(198144L); i0+=static_cast<long>(16L))
        {
            auto tmp0 = at::vec::load_uint8_as_float(in_ptr0 + static_cast<long>(i0));
            auto tmp1 = (tmp0);
            auto tmp2 = at::vec::Vectorized<float>(static_cast<float>(1.0));
            auto tmp3 = tmp1 - tmp2;
            auto tmp4 = at::vec::Vectorized<float>(static_cast<float>(0.10000000149011612));
            auto tmp5 = tmp3 * tmp4;
            auto tmp6 = at::vec::clamp_min(tmp5, decltype(tmp5)(0));
            auto tmp7 = tmp4.reciprocal();
            auto tmp8 = tmp7 * tmp2;
            auto tmp9 = tmp6 * tmp8;
            auto tmp10 = tmp9.round();
            auto tmp11 = tmp10 + tmp2;
            auto tmp12 = at::vec::Vectorized<float>(static_cast<float>(0.0));
            auto tmp13 = at::vec::maximum(tmp11, tmp12);
            auto tmp14 = at::vec::Vectorized<float>(static_cast<float>(255.0));
            auto tmp15 = at::vec::minimum(tmp13, tmp14);
            auto tmp16 = (tmp15);
            at::vec::store_float_as_uint8(tmp16, out_ptr0 + static_cast<long>(i0));
        }
        #pragma omp simd simdlen(8)
        for(long i0=static_cast<long>(198144L); i0<static_cast<long>(198147L); i0+=static_cast<long>(1L))
        {
            auto tmp0 = in_ptr0[static_cast<long>(i0)];
            auto tmp1 = static_cast<float>(tmp0);
            auto tmp2 = static_cast<float>(1.0);
            auto tmp3 = tmp1 - tmp2;
            auto tmp4 = static_cast<float>(0.10000000149011612);
            auto tmp5 = decltype(tmp3)(tmp3 * tmp4);
            auto tmp6 = tmp5 * (tmp5>0);
            auto tmp7 = 1 / tmp4;
            auto tmp8 = decltype(tmp7)(tmp7 * tmp2);
            auto tmp9 = decltype(tmp6)(tmp6 * tmp8);
            auto tmp10 = std::nearbyint(tmp9);
            auto tmp11 = tmp10 + tmp2;
            auto tmp12 = static_cast<float>(0.0);
            auto tmp13 = max_propagate_nan(tmp11, tmp12);
            auto tmp14 = static_cast<float>(255.0);
            auto tmp15 = min_propagate_nan(tmp13, tmp14);
            auto tmp16 = static_cast<unsigned char>(tmp15);
            out_ptr0[static_cast<long>(i0)] = tmp16;
        }
    }
}
''')

```

## Additional context

This RFC focus on x86 CPU device. As for the other devices, different quantizers must be added per device into the frontend. Meanwhile, lowering the quantized graph via inductor compiler should be similar.

cc @jerryzh168 @jianyuh @raghuramank100 @jamesr66a @vkuzo @jgong5 @Xia-Weiwen @ezyang @msaroufim @wconstab @bdhirsh @anijain2305