Add UInt16 support to Q and DQ

Differential Revision: D65899080

Pull Request resolved: #6891

Author: JacobSzwejbka

kernels/quantized/cpu/op_dequantize.cpp

@@ -39,6 +39,7 @@ void check_dequantize_per_tensor_args(
       input.scalar_type() == ScalarType::Byte ||
           input.scalar_type() == ScalarType::Char ||
           input.scalar_type() == ScalarType::Bits16 ||
+          input.scalar_type() == ScalarType::UInt16 ||
           input.scalar_type() == ScalarType::Short ||
           input.scalar_type() == ScalarType::Int,
       "input.scalar_type() %" PRId8 " is not supported:",
@@ -120,6 +121,7 @@ Tensor& dequantize_per_tensor_out(
   switch (input.scalar_type()) {
     ET_FORALL_INT_TYPES(CALCULATE_INT_TYPE);
     CALCULATE_INT_TYPE(uint16_t, Bits16);
+    CALCULATE_INT_TYPE(uint16_t, UInt16);
     default:
       ET_CHECK_MSG(
           false,
@@ -315,6 +317,7 @@ Tensor& dequantize_per_channel_out(
   switch (input.scalar_type()) {
     ET_FORALL_INT_TYPES(CALCULATE_FLOAT_TYPE);
     CALCULATE_INT_TYPE(uint16_t, Bits16);
+    CALCULATE_INT_TYPE(uint16_t, UInt16);
     default:
       ET_CHECK_MSG(
           false,

kernels/quantized/cpu/op_quantize.cpp

@@ -57,7 +57,7 @@ void check_quantize_per_tensor_args(
         static_cast<int32_t>(std::numeric_limits<int8_t>::min());
     quant_max_upper_bound =
         static_cast<int32_t>(std::numeric_limits<int8_t>::max());
-  } else if (dtype == ScalarType::Bits16) {
+  } else if (dtype == ScalarType::Bits16 || dtype == ScalarType::UInt16) {
     quant_min_lower_bound = std::numeric_limits<uint16_t>::min();
     quant_max_upper_bound = std::numeric_limits<uint16_t>::max();
   } else if (dtype == ScalarType::Short) {
@@ -139,6 +139,7 @@ Tensor& quantize_per_tensor_out(
     switch (out.scalar_type()) {                         \
       ET_FORALL_INT_TYPES_WITH(IN_CTYPE, QUANTIZE_IMPL); \
       QUANTIZE_IMPL(IN_CTYPE, uint16_t, Bits16)          \
+      QUANTIZE_IMPL(IN_CTYPE, uint16_t, UInt16)          \
       default:                                           \
         ET_CHECK_MSG(                                    \
             false,                                       \
@@ -334,6 +335,7 @@ Tensor& quantize_per_channel_out(
     switch (out.scalar_type()) {                         \
       ET_FORALL_INT_TYPES_WITH(CTYPE_IN, QUANTIZE_IMPL); \
       QUANTIZE_IMPL(CTYPE_IN, uint16_t, Bits16)          \
+      QUANTIZE_IMPL(CTYPE_IN, uint16_t, UInt16)          \
       default:                                           \
         ET_CHECK_MSG(                                    \
             false,                                       \

kernels/quantized/test/op_dequantize_test.cpp

@@ -63,6 +63,7 @@ TEST(OpDequantizeOutTest, AllDtypesSupported) {
   test_dtype<ScalarType::Char>();
   test_dtype<ScalarType::Short>();
   test_dtype<ScalarType::Bits16>();
+  test_dtype<ScalarType::UInt16>();
   test_dtype<ScalarType::Int>();
 }
 

kernels/quantized/test/op_quantize_test.cpp

@@ -54,6 +54,7 @@ TEST(OpQuantizeOutTest, AllDtypesSupported) {
   test_dtype<ScalarType::Char>();
   test_dtype<ScalarType::Short>();
   test_dtype<ScalarType::Bits16>();
+  test_dtype<ScalarType::UInt16>();
   test_dtype<ScalarType::Int>();
 }
 

runtime/core/exec_aten/testing_util/tensor_factory.h

@@ -650,6 +650,13 @@ struct ScalarTypeToCppTypeWrapper<torch::executor::ScalarType::Bits16> {
   using ctype = uint16_t;
 };
 
+// Use a C type of `uint16_t` instead of `UInt16` to simplify code reuse when
+// testing multiple integer types.
+template <>
+struct ScalarTypeToCppTypeWrapper<torch::executor::ScalarType::UInt16> {
+  using ctype = uint16_t;
+};
+
 // To allow implicit conversion between simple types to `ctype`
 #define SPECIALIZE_ScalarTypeToCppTypeWrapper(CTYPE, DTYPE)               \
   template <>                                                             \