Add support for uint16 in quant and dequant kernels

Differential Revision: D65370235

Pull Request resolved: #6724

Author: dulinriley

backends/cadence/hifi/kernels/kernels.cpp

@@ -165,6 +165,7 @@ void requantize(
 typed_quantize_val(int8_t);
 typed_quantize_val(uint8_t);
 typed_quantize_val(int16_t);
+typed_quantize_val(uint16_t);
 #undef typed_quantize_val
 
 #define typed_quantize_vec(dtype)  \
@@ -177,6 +178,7 @@ typed_quantize_val(int16_t);
 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
 
@@ -186,6 +188,7 @@ typed_quantize_vec(int32_t);
 typed_dequantize_val(int8_t);
 typed_dequantize_val(uint8_t);
 typed_dequantize_val(int16_t);
+typed_dequantize_val(uint16_t);
 #undef typed_dequantize_val
 
 #define typed_dequantize_vec(dtype) \
@@ -198,6 +201,7 @@ typed_dequantize_val(int16_t);
 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
 

backends/cadence/hifi/operators/dequantize_per_tensor.cpp

@@ -41,6 +41,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::Bits16) {
+    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);

backends/cadence/hifi/operators/quantize_per_tensor.cpp

@@ -44,6 +44,10 @@ void quantize_per_tensor_out(
     int16_t* out_data = out.mutable_data_ptr<int16_t>();
     cadence::impl::HiFi::kernels::quantize<int16_t>(
         out_data, input_data, 1. / scale, zero_point, numel);
+  } else if (out.scalar_type() == ScalarType::Bits16) {
+    uint16_t* out_data = out.mutable_data_ptr<uint16_t>();
+    cadence::impl::HiFi::kernels::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>();
     cadence::impl::HiFi::kernels::quantize<int32_t>(

backends/cadence/reference/kernels/kernels.cpp

@@ -65,6 +65,7 @@ void dequantize(
 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
 
@@ -78,6 +79,7 @@ typed_quantize_val(int32_t);
 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
 
@@ -86,6 +88,7 @@ typed_quantize_vec(int32_t);
 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
 
@@ -99,6 +102,7 @@ typed_dequantize_val(int32_t);
 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
 

backends/cadence/reference/operators/dequantize_per_tensor.cpp

@@ -37,6 +37,14 @@ void dequantize_per_tensor_out(
     const int8_t* input_data = input.const_data_ptr<int8_t>();
     impl::reference::kernels::dequantize<int8_t>(
         out_data, input_data, scale, zero_point, numel);
+  } else if (input.scalar_type() == ScalarType::Bits16) {
+    const uint16_t* input_data = input.const_data_ptr<uint16_t>();
+    impl::reference::kernels::dequantize<uint16_t>(
+        out_data, input_data, scale, zero_point, numel);
+  } else if (input.scalar_type() == ScalarType::Short) {
+    const int16_t* input_data = input.const_data_ptr<int16_t>();
+    impl::reference::kernels::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>();
     impl::reference::kernels::dequantize<int32_t>(

backends/cadence/reference/operators/quantize_per_tensor.cpp

@@ -39,6 +39,14 @@ void quantize_per_tensor_out(
     int8_t* out_data = out.mutable_data_ptr<int8_t>();
     impl::reference::kernels::quantize<int8_t>(
         out_data, input_data, 1. / scale, zero_point, numel);
+  } else if (out.scalar_type() == ScalarType::Bits16) {
+    uint16_t* out_data = out.mutable_data_ptr<uint16_t>();
+    impl::reference::kernels::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>();
+    impl::reference::kernels::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>();
     impl::reference::kernels::quantize<int32_t>(