Add int32 quant/dequant back (#14269)

Summary:
Pull Request resolved: #14269

Now that the previous diff exists, we can add the int32 case back without adding to the code size of deployed models.

Reviewed By: hsharma35

Differential Revision: D82282481

Author: mcremon-meta

backends/cadence/aot/functions.yaml

@@ -208,6 +208,12 @@
     - arg_meta: null
       kernel_name: impl::generic::quantize_per_tensor_asym16u_out
 
+- func: cadence::quantize_per_tensor_asym32s.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)
+  variants: function
+  kernels:
+    - arg_meta: null
+      kernel_name: impl::generic::quantize_per_tensor_asym32s_out
+
 - func: cadence::dequantize_per_tensor.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)
   variants: function
   kernels:
@@ -238,6 +244,12 @@
     - arg_meta: null
       kernel_name: impl::generic::dequantize_per_tensor_asym16u_out
 
+- func: cadence::dequantize_per_tensor_asym32s.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)
+  variants: function
+  kernels:
+    - arg_meta: null
+      kernel_name: impl::generic::dequantize_per_tensor_asym32s_out
+
 - func: cadence::quantized_conv_nchw.out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, Tensor weight_zero_point, Tensor bias_scale, float out_scale, int out_zero_point, Tensor out_multiplier, Tensor out_shift, *, Tensor(a!) out) -> Tensor(a!)
   kernels:
     - arg_meta: null

backends/cadence/aot/functions_hifi.yaml

@@ -308,6 +308,11 @@
     - arg_meta: null
       kernel_name: impl::HiFi::quantize_per_tensor_asym16s_out
 
+- func: cadence::quantize_per_tensor_asym32s.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)
+  variants: function
+  kernels:
+    - arg_meta: null
+      kernel_name: impl::HiFi::quantize_per_tensor_asym32s_out
 
 - func: cadence::dequantize_per_tensor.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)
   variants: function
@@ -339,6 +344,12 @@
     - arg_meta: null
       kernel_name: impl::HiFi::dequantize_per_tensor_asym16u_out
 
+- func: cadence::dequantize_per_tensor_asym32s.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)
+  variants: function
+  kernels:
+    - arg_meta: null
+      kernel_name: impl::HiFi::dequantize_per_tensor_asym16s_out
+
 - func: cadence::quantized_conv_nchw.out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, Tensor weight_zero_point, Tensor bias_scale, float out_scale, int out_zero_point, Tensor out_multiplier, Tensor out_shift, *, Tensor(a!) out) -> Tensor(a!)
   kernels:
     - arg_meta: null

backends/cadence/aot/ops_registrations.py

@@ -56,6 +56,13 @@
     "quantize_per_tensor_asym16u.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)"
 )
 
+lib.define(
+    "quantize_per_tensor_asym32s(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype) -> (Tensor Z)"
+)
+lib.define(
+    "quantize_per_tensor_asym32s.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)"
+)
+
 lib.define(
     "dequantize_per_tensor(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype) -> (Tensor Z)"
 )
@@ -87,6 +94,13 @@
     "dequantize_per_tensor_asym16u.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)"
 )
 
+lib.define(
+    "dequantize_per_tensor_asym32s(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype) -> (Tensor Z)"
+)
+lib.define(
+    "dequantize_per_tensor_asym32s.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)"
+)
+
 lib.define(
     "quantized_layer_norm(Tensor X, Tensor X_scale, Tensor X_zero_point, int[] normalized_shape, Tensor weight, Tensor bias, float eps, float output_scale, int output_zero_point) -> (Tensor Y)"
 )
@@ -617,6 +631,18 @@ def quantize_per_tensor_asym16u_meta(
     return input.new_empty(input.size(), dtype=dtype)
 
 
+@register_fake("cadence::quantize_per_tensor_asym32s")
+def quantize_per_tensor_asym32s_meta(
+    input: torch.Tensor,
+    scale: float,
+    zero_point: int,
+    quant_min: int,
+    quant_max: int,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    return input.new_empty(input.size(), dtype=dtype)
+
+
 @register_fake("cadence::dequantize_per_tensor")
 def dequantize_per_tensor_meta(
     input: torch.Tensor,
@@ -677,6 +703,18 @@ def dequantize_per_tensor_asym16u_meta(
     return input.new_empty(input.size(), dtype=torch.float)
 
 
+@register_fake("cadence::dequantize_per_tensor_asym32s")
+def dequantize_per_tensor_asym32s_meta(
+    input: torch.Tensor,
+    scale: float,
+    zero_point: int,
+    quant_min: int,
+    quant_max: int,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    return input.new_empty(input.size(), dtype=torch.float)
+
+
 @register_fake("cadence::quantized_add")
 def quantized_add_meta(
     X: torch.Tensor,

backends/cadence/aot/type_dispatch.py

@@ -108,6 +108,7 @@ class CompileTimeTypeDispatchPass(ExportPass):
                 (torch.uint8,): "asym8u",
                 (torch.int16,): "asym16s",
                 (torch.uint16,): "asym16s",
+                (torch.int32,): "asym32s",
             },
             variant="default",
             is_quant_op=True,
@@ -119,6 +120,7 @@ class CompileTimeTypeDispatchPass(ExportPass):
                 (torch.uint8,): "asym8u",
                 (torch.int16,): "asym16s",
                 (torch.uint16,): "asym16s",
+                (torch.int32,): "asym32s",
             },
             variant="default",
         ),

backends/cadence/generic/kernels/kernels.cpp

@@ -73,6 +73,7 @@ typed_quantize_val(int8_t);
 typed_quantize_val(uint8_t);
 typed_quantize_val(int16_t);
 typed_quantize_val(uint16_t);
+typed_quantize_val(int32_t);
 #undef typed_quantize_val
 
 #define typed_quantize_vec(dtype)  \
@@ -86,6 +87,7 @@ typed_quantize_vec(int8_t);
 typed_quantize_vec(uint8_t);
 typed_quantize_vec(int16_t);
 typed_quantize_vec(uint16_t);
+typed_quantize_vec(int32_t);
 #undef typed_quantize_vec
 
 #define typed_dequantize_val(dtype) \
@@ -94,6 +96,7 @@ typed_dequantize_val(int8_t);
 typed_dequantize_val(uint8_t);
 typed_dequantize_val(int16_t);
 typed_dequantize_val(uint16_t);
+typed_dequantize_val(int32_t);
 #undef typed_dequantize_val
 
 #define typed_dequantize_vec(dtype) \
@@ -107,6 +110,7 @@ typed_dequantize_vec(int8_t);
 typed_dequantize_vec(uint8_t);
 typed_dequantize_vec(int16_t);
 typed_dequantize_vec(uint16_t);
+typed_dequantize_vec(int32_t);
 #undef typed_dequantize_vec
 
 } // namespace kernels

backends/cadence/generic/operators/dequantize_per_tensor.cpp

@@ -44,6 +44,9 @@ void dequantize_per_tensor_out(
   } else if (input.scalar_type() == ScalarType::Short) {
     const int16_t* input_data = input.const_data_ptr<int16_t>();
     dequantize<int16_t>(out_data, input_data, scale, zero_point, numel);
+  } else if (input.scalar_type() == ScalarType::Int) {
+    const int32_t* input_data = input.const_data_ptr<int32_t>();
+    dequantize<int32_t>(out_data, input_data, scale, zero_point, numel);
   } else {
     ET_CHECK_MSG(
         false,
@@ -112,6 +115,21 @@ void dequantize_per_tensor_asym16u_out(
   dequantize<uint16_t>(out_data, input_data, scale, zero_point, numel);
 }
 
+void dequantize_per_tensor_asym32s_out(
+    KernelRuntimeContext& context,
+    const Tensor& input,
+    double scale,
+    int64_t zero_point,
+    int64_t quant_min,
+    int64_t quant_max,
+    ScalarType dtype,
+    Tensor& out) {
+  float* out_data = out.mutable_data_ptr<float>();
+  size_t numel = out.numel();
+  const int32_t* input_data = input.const_data_ptr<int32_t>();
+  dequantize<int32_t>(out_data, input_data, scale, zero_point, numel);
+}
+
 } // namespace native
 } // namespace generic
 } // namespace impl

backends/cadence/generic/operators/quantize_per_tensor.cpp

@@ -34,22 +34,22 @@ void quantize_per_tensor_out(
 
   if (out.scalar_type() == ScalarType::Byte) {
     uint8_t* out_data = out.mutable_data_ptr<uint8_t>();
-    quantize<uint8_t>(
-        out_data, input_data, 1. / scale, zero_point, numel);
+    quantize<uint8_t>(out_data, input_data, 1. / scale, zero_point, numel);
   } else if (out.scalar_type() == ScalarType::Char) {
     int8_t* out_data = out.mutable_data_ptr<int8_t>();
-    quantize<int8_t>(
-        out_data, input_data, 1. / scale, zero_point, numel);
+    quantize<int8_t>(out_data, input_data, 1. / scale, zero_point, numel);
   } else if (
       out.scalar_type() == ScalarType::Bits16 ||
       out.scalar_type() == ScalarType::UInt16) {
     uint16_t* out_data = out.mutable_data_ptr<uint16_t>();
-    quantize<uint16_t>(
-        out_data, input_data, 1. / scale, zero_point, numel);
+    quantize<uint16_t>(out_data, input_data, 1. / scale, zero_point, numel);
   } else if (out.scalar_type() == ScalarType::Short) {
     int16_t* out_data = out.mutable_data_ptr<int16_t>();
     quantize<int16_t>(
         out_data, input_data, 1. / scale, zero_point, numel);
+  } else if (out.scalar_type() == ScalarType::Int) {
+    int32_t* out_data = out.mutable_data_ptr<int32_t>();
+    quantize<int32_t>(out_data, input_data, 1. / scale, zero_point, numel);
   } else {
     ET_CHECK_MSG(
         false,
@@ -70,8 +70,7 @@ void quantize_per_tensor_asym8s_out(
   const float* input_data = input.const_data_ptr<float>();
   size_t numel = out.numel();
   int8_t* out_data = out.mutable_data_ptr<int8_t>();
-  impl::generic::kernels::quantize<int8_t>(
-      out_data, input_data, 1. / scale, zero_point, numel);
+  quantize<int8_t>(out_data, input_data, 1. / scale, zero_point, numel);
 }
 
 void quantize_per_tensor_asym8u_out(
@@ -86,8 +85,7 @@ void quantize_per_tensor_asym8u_out(
   const float* input_data = input.const_data_ptr<float>();
   size_t numel = out.numel();
   uint8_t* out_data = out.mutable_data_ptr<uint8_t>();
-  impl::generic::kernels::quantize<uint8_t>(
-      out_data, input_data, 1. / scale, zero_point, numel);
+  quantize<uint8_t>(out_data, input_data, 1. / scale, zero_point, numel);
 }
 
 void quantize_per_tensor_asym16s_out(
@@ -102,8 +100,7 @@ void quantize_per_tensor_asym16s_out(
   const float* input_data = input.const_data_ptr<float>();
   size_t numel = out.numel();
   int16_t* out_data = out.mutable_data_ptr<int16_t>();
-  impl::generic::kernels::quantize<int16_t>(
-      out_data, input_data, 1. / scale, zero_point, numel);
+  quantize<int16_t>(out_data, input_data, 1. / scale, zero_point, numel);
 }
 
 void quantize_per_tensor_asym16u_out(
@@ -118,8 +115,22 @@ void quantize_per_tensor_asym16u_out(
   const float* input_data = input.const_data_ptr<float>();
   size_t numel = out.numel();
   uint16_t* out_data = out.mutable_data_ptr<uint16_t>();
-  impl::generic::kernels::quantize<uint16_t>(
-      out_data, input_data, 1. / scale, zero_point, numel);
+  quantize<uint16_t>(out_data, input_data, 1. / scale, zero_point, numel);
+}
+
+void quantize_per_tensor_asym32s_out(
+    KernelRuntimeContext& context,
+    const Tensor& input,
+    double scale,
+    int64_t zero_point,
+    int64_t quant_min,
+    int64_t quant_max,
+    ScalarType dtype,
+    Tensor& out) {
+  const float* input_data = input.const_data_ptr<float>();
+  size_t numel = out.numel();
+  int32_t* out_data = out.mutable_data_ptr<int32_t>();
+  quantize<int32_t>(out_data, input_data, 1. / scale, zero_point, numel);
 }
 
 } // namespace native

backends/cadence/hifi/kernels/kernels.cpp

@@ -135,6 +135,7 @@ typed_quantize_val(int8_t);
 typed_quantize_val(uint8_t);
 typed_quantize_val(int16_t);
 typed_quantize_val(uint16_t);
+typed_quantize_val(int32_t);
 #undef typed_quantize_val
 
 #define typed_quantize_vec(dtype)  \
@@ -158,6 +159,7 @@ typed_dequantize_val(int8_t);
 typed_dequantize_val(uint8_t);
 typed_dequantize_val(int16_t);
 typed_dequantize_val(uint16_t);
+typed_dequantize_val(int32_t);
 #undef typed_dequantize_val
 
 #define typed_dequantize_vec(dtype) \

backends/cadence/hifi/operators/op_dequantize_per_tensor.cpp

@@ -45,6 +45,9 @@ void dequantize_per_tensor_out(
       input.scalar_type() == ScalarType::UInt16) {
     const uint16_t* input_data = input.const_data_ptr<uint16_t>();
     dequantize<uint16_t>(out_data, input_data, scale, zero_point, numel);
+  } else if (input.scalar_type() == ScalarType::Int) {
+    const int32_t* input_data = input.const_data_ptr<int32_t>();
+    dequantize<int32_t>(out_data, input_data, scale, zero_point, numel);
   } else {
     ET_CHECK_MSG(
         false,
@@ -98,6 +101,21 @@ void dequantize_per_tensor_asym16u_out(
   dequantize<uint16_t>(out_data, input_data, scale, zero_point, numel);
 }
 
+void dequantize_per_tensor_asym32s_out(
+    KernelRuntimeContext& context,
+    const Tensor& input,
+    double scale,
+    int64_t zero_point,
+    int64_t quant_min,
+    int64_t quant_max,
+    ScalarType dtype,
+    Tensor& out) {
+  float* out_data = out.mutable_data_ptr<float>();
+  size_t numel = out.numel();
+  const int32_t* input_data = input.const_data_ptr<int32_t>();
+  dequantize<int32_t>(out_data, input_data, scale, zero_point, numel);
+}
+
 } // namespace native
 } // namespace HiFi
 } // namespace impl

backends/cadence/hifi/operators/op_quantize_per_tensor.cpp

@@ -105,6 +105,9 @@ void quantize_per_tensor_out(
       out.scalar_type() == ScalarType::UInt16) {
     uint16_t* out_data = out.mutable_data_ptr<uint16_t>();
     quantize<uint16_t>(out_data, input_data, 1. / scale, zero_point, numel);
+  } else if (out.scalar_type() == ScalarType::Int) {
+    int32_t* out_data = out.mutable_data_ptr<int32_t>();
+    quantize<int32_t>(out_data, input_data, 1. / scale, zero_point, numel);
   } else {
     ET_KERNEL_CHECK_MSG(
         ctx,
@@ -161,6 +164,21 @@ void quantize_per_tensor_asym16u_out(
   quantize<uint16_t>(out_data, input_data, 1. / scale, zero_point, numel);
 }
 
+void quantize_per_tensor_asym32s_out(
+    KernelRuntimeContext& context,
+    const Tensor& input,
+    double scale,
+    int64_t zero_point,
+    int64_t quant_min,
+    int64_t quant_max,
+    ScalarType dtype,
+    Tensor& out) {
+  const float* input_data = input.const_data_ptr<float>();
+  size_t numel = out.numel();
+  int32_t* out_data = out.mutable_data_ptr<int32_t>();
+  quantize<int32_t>(out_data, input_data, 1. / scale, zero_point, numel);
+}
+
 } // namespace native
 } // namespace HiFi
 } // namespace impl