Dynamic dispatch for optimized quantized op kernels (#25545)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/25545

This re-uses the infrastructure from ATen/native/cpu, which compiles kernels multiple times for different instruction sets and dispatches dynamically based on the CPU's capability flags at runtime. This ensures we use the most optimal quantized kernel for the given machine

Test Plan: Imported from OSS

Differential Revision: D17166369

Pulled By: jamesr66a

fbshipit-source-id: 8c8393f99365e1408819bbaf254c1b5734a34b70

aten/src/ATen/cpu/vec256/vec256_qint.h

@@ -294,8 +294,8 @@ struct Vec256<c10::quint8> {
         Vec256<c10::quint8> zero_point,
         Vec256<c10::quint8> q_six) {
 #ifdef __AVX2__
-      return _mm256_min_epi8(
-          _mm256_max_epi8(vals, zero_point.vals), q_six.vals);
+      return _mm256_min_epu8(
+          _mm256_max_epu8(vals, zero_point.vals), q_six.vals);
 #else
       // Pray the compiler can autovectorize this
       uint8_t int_vals[size()];
@@ -405,8 +405,8 @@ struct Vec256<c10::qint32> {
         Vec256<c10::qint32> zero_point,
         Vec256<c10::qint32> q_six) {
 #ifdef __AVX2__
-      return _mm256_min_epi8(
-          _mm256_max_epi8(vals, zero_point.vals), q_six.vals);
+      return _mm256_min_epi32(
+          _mm256_max_epi32(vals, zero_point.vals), q_six.vals);
 #else
       // Pray the compiler can autovectorize this
       int32_t int_vals[size()];

aten/src/ATen/native/quantized/cpu/kernels/QuantizedOpKernels.cpp

@@ -0,0 +1,128 @@
+#include <ATen/ATen.h>
+#include <ATen/Dispatch.h>
+#include <ATen/native/TensorIterator.h>
+#include <ATen/native/cpu/Loops.h>
+#include <ATen/native/quantized/cpu/quantized_ops.h>
+
+namespace at {
+namespace native {
+namespace {
+
+// ****************** HEY YOU! YES YOU! Read this! ********************
+//
+// Please read the README.md in this directory before editing this file
+
+void qrelu_kernel(const Tensor& qx, Tensor& qy) {
+  const auto zero_point = qx.q_zero_point();
+  AT_DISPATCH_QINT_TYPES(qx.scalar_type(), "qrelu", [&]() {
+    qy = at::_empty_affine_quantized(
+        qx.sizes(),
+        at::device(kCPU).dtype(SCALAR_TYPE),
+        qx.q_scale(),
+        qx.q_zero_point(),
+        qx.suggest_memory_format());
+    using Vec = Vec256<scalar_t>;
+    auto zero_point_vec = Vec(scalar_t(zero_point));
+    auto iter = TensorIterator::unary_op(qy, qx);
+    cpu_kernel_vec(
+        iter,
+        [&](scalar_t value) -> scalar_t {
+          return scalar_t(std::max<underlying_t>(value.val_, zero_point));
+        },
+        [&](Vec value) -> Vec { return value.relu(zero_point_vec); });
+  });
+}
+
+void qrelu6_kernel(const Tensor& qx, Tensor& qy) {
+  const auto zero_point = qx.q_zero_point();
+  AT_DISPATCH_QINT_TYPES(qx.scalar_type(), "qrelu6", [&]() {
+    qy = at::_empty_affine_quantized(
+        qx.sizes(),
+        at::device(kCPU).dtype(SCALAR_TYPE),
+        qx.q_scale(),
+        qx.q_zero_point(),
+        qx.suggest_memory_format());
+    using Vec = Vec256<scalar_t>;
+    auto iter = TensorIterator::unary_op(qy, qx);
+    scalar_t six =
+        at::quantize_val<scalar_t>(qx.q_scale(), qx.q_zero_point(), 6.0);
+    auto zero_point_vec = Vec(scalar_t(zero_point));
+    auto six_vec = Vec(six);
+    cpu_kernel_vec(
+        iter,
+        [&](scalar_t value) -> scalar_t {
+          underlying_t relu_val =
+              std::max<underlying_t>(value.val_, zero_point);
+          return scalar_t(std::min<underlying_t>(relu_val, six.val_));
+        },
+        [&](Vec val) -> Vec { return val.relu6(zero_point_vec, six_vec); });
+  });
+}
+
+// Note: out is assumed to be the same size as self and other.
+// Note: Addition is only supported when self, other, out are of the same dtype.
+template <bool ReLUFused = false>
+void qadd_kernel(Tensor& out, const Tensor& self, const Tensor& other) {
+  int64_t zero_point = out.q_zero_point();
+  double scale = out.q_scale();
+  int64_t self_zero_point = self.q_zero_point();
+  double self_scale = self.q_scale();
+  int64_t other_zero_point = other.q_zero_point();
+  double other_scale = other.q_scale();
+
+  // Broadcast out the parameters here to amortize out that cost across
+  // loop iterations.
+  // TODO: we can optimize dequantization by doing a premultiplication
+  // of the zero point by scale and doing FMA on scale*x_q - (scale*zero_point)
+  auto self_zero_point_vec = Vec256<float>((float)self_zero_point);
+  auto self_scale_vec = Vec256<float>(self_scale);
+  auto other_zero_point_vec = Vec256<float>((float)other_zero_point);
+  auto other_scale_vec = Vec256<float>(other_scale);
+
+  auto iter = TensorIterator::binary_op(out, self, other);
+
+  AT_DISPATCH_QINT_TYPES(out.scalar_type(), "qadd", [&]() {
+    using Vec = Vec256<scalar_t>;
+    cpu_kernel_vec(
+        iter,
+        [&](scalar_t a, scalar_t b) -> scalar_t {
+          const auto da = at::dequantize_val(self_scale, self_zero_point, a);
+          const auto db = at::dequantize_val(other_scale, other_zero_point, b);
+          float c = da + db;
+          if (ReLUFused) {
+            c = std::max<float>(c, 0.0);
+          }
+          return at::quantize_val<scalar_t>(scale, zero_point, c);
+        },
+        [&](Vec a, Vec b) -> Vec {
+          const auto da = a.dequantize(self_scale_vec, self_zero_point_vec);
+          const auto db = b.dequantize(other_scale_vec, other_zero_point_vec);
+          Vec::float_vec_return_type retvals;
+          for (int i = 0; i < Vec::float_num_vecs(); ++i) {
+            auto c = da[i] + db[i];
+            if (ReLUFused) {
+              c = vec256::maximum(c, Vec256<float>(0.0f));
+            }
+            retvals[i] = c;
+          }
+          // TODO: fbgemm::Quantize doesn't support taking in the
+          // pre-broadcasted parameters. We might be able to save some cycles by
+          // enabling that in the API.
+          // TODO: specialize fbgemm::Quantize for a single vector and make it
+          // inlineable. This could help with interleaving as suggested by the
+          // TensorIterator implementations
+          auto rv = Vec::quantize(retvals, scale, zero_point);
+          return rv;
+        });
+  });
+}
+
+} // namespace
+
+REGISTER_DISPATCH(qrelu_stub, &qrelu_kernel);
+REGISTER_DISPATCH(qrelu6_stub, &qrelu6_kernel);
+REGISTER_DISPATCH(qadd_relu_stub, &qadd_kernel<true>);
+REGISTER_DISPATCH(qadd_stub, &qadd_kernel<false>);
+
+} // namespace native
+} // namespace at

aten/src/ATen/native/quantized/cpu/kernels/README.md

@@ -0,0 +1,16 @@
+ The files in this directory are compiled multiple times for different CPU vector instruction
+ sets (e.g. AVX, AVX2). The purpose of putting code in this directory is to make
+ sure we can generate the optimal code for a given processor's vector
+ capabilities. Much of this is done via preprocessor guards in vec256_qint.h.
+
+ The considerations for code written in this directory include:
+  - Keep code in this directory to a minimum, since we're compiling it several
+    times.
+  - All code in this file should go through the DECLARE_DISPATCH,
+    DEFINE_DISPATCH, and REGISTER_DISPATCH mechanism to ensure the correct
+    runtime dispatch occurs.
+  - THE CODE MUST RESIDE IN THE ANONYMOUS NAMESPACE. FAILURE TO ENSURE THIS
+    IS THE CASE CAN LEAD TO HARD-TO-DEBUG ODR VIOLATIONS.
+  - **Make sure different variants of the code (AVX, AVX2) are tested!**
+    There are build variants that do things like have NO AVX and NO AVX2 in
+    CI. Make sure they work!

aten/src/ATen/native/quantized/cpu/qadd.cpp

@@ -4,11 +4,16 @@
 #include <ATen/native/TensorIterator.h>
 #include <ATen/native/cpu/Loops.h>
 #include <ATen/quantized/Quantizer.h>
+#include <ATen/native/quantized/cpu/quantized_ops.h>
 
 #include <algorithm>
 
 namespace at {
 namespace native {
+
+DEFINE_DISPATCH(qadd_relu_stub);
+DEFINE_DISPATCH(qadd_stub);
+
 namespace {
 
 inline void check_inputs(const Tensor& qa, const Tensor& qb) {
@@ -28,56 +33,11 @@ inline void check_inputs(const Tensor& qa, const Tensor& qb) {
 // Note: Addition is only supported when self, other, out are of the same dtype.
 template <bool ReLUFused = false>
 Tensor _add_out(Tensor& out, const Tensor& self, const Tensor& other) {
-  int64_t zero_point = out.q_zero_point();
-  double scale = out.q_scale();
-  int64_t self_zero_point = self.q_zero_point();
-  double self_scale = self.q_scale();
-  int64_t other_zero_point = other.q_zero_point();
-  double other_scale = other.q_scale();
-
-  // Broadcast out the parameters here to amortize out that cost across
-  // loop iterations.
-  // TODO: we can optimize dequantization by doing a premultiplication
-  // of the zero point by scale and doing FMA on scale*x_q - (scale*zero_point)
-  auto self_zero_point_vec = Vec256<float>((float)self_zero_point);
-  auto self_scale_vec = Vec256<float>(self_scale);
-  auto other_zero_point_vec = Vec256<float>((float)other_zero_point);
-  auto other_scale_vec = Vec256<float>(other_scale);
-
-  auto iter = TensorIterator::binary_op(out, self, other);
-
-  AT_DISPATCH_QINT_TYPES(out.scalar_type(), "qadd", [&]() {
-    using Vec = Vec256<scalar_t>;
-    cpu_kernel_vec(iter, [&](scalar_t a, scalar_t b) -> scalar_t {
-      const auto da = at::dequantize_val(self_scale, self_zero_point, a);
-      const auto db = at::dequantize_val(other_scale, other_zero_point, b);
-      float c = da + db;
-      if (ReLUFused) {
-        c = std::max<float>(c, 0.0);
-      }
-      return at::quantize_val<scalar_t>(scale, zero_point, c);
-    },
-    [&](Vec a, Vec b) -> Vec {
-      const auto da = a.dequantize(self_scale_vec, self_zero_point_vec);
-      const auto db = b.dequantize(other_scale_vec, other_zero_point_vec);
-      Vec::float_vec_return_type retvals;
-      for (int i = 0; i < Vec::float_num_vecs(); ++i) {
-        auto c = da[i] + db[i];
-        if (ReLUFused) {
-          c = vec256::maximum(c, Vec256<float>(0.0f));
-        }
-        retvals[i] = c;
-      }
-      // TODO: fbgemm::Quantize doesn't support taking in the pre-broadcasted
-      // parameters. We might be able to save some cycles by enabling that
-      // in the API.
-      // TODO: specialize fbgemm::Quantize for a single vector and make it
-      // inlineable. This could help with interleaving as suggested by the
-      // TensorIterator implementations
-      auto rv = Vec::quantize(retvals, scale, zero_point);
-      return rv;
-    });
-  });
+  if (ReLUFused) {
+    qadd_relu_stub(self.device().type(), out, self, other);
+  } else {
+    qadd_stub(self.device().type(), out, self, other);
+  }
   return out;
 }
 

aten/src/ATen/native/quantized/cpu/qrelu.cpp

@@ -4,31 +4,19 @@
 #include <ATen/native/TensorIterator.h>
 #include <ATen/native/cpu/Loops.h>
 #include <ATen/quantized/Quantizer.h>
+#include <ATen/native/quantized/cpu/quantized_ops.h>
 
 #include <algorithm>
 
 namespace at {
 namespace native {
+
+DEFINE_DISPATCH(qrelu_stub);
+DEFINE_DISPATCH(qrelu6_stub);
+
 Tensor quantized_relu(const Tensor& qx) {
   Tensor qy;
-  const auto zero_point = qx.q_zero_point();
-  AT_DISPATCH_QINT_TYPES(qx.scalar_type(), "qrelu", [&]() {
-    qy = at::_empty_affine_quantized(
-        qx.sizes(),
-        at::device(kCPU).dtype(SCALAR_TYPE),
-        qx.q_scale(),
-        qx.q_zero_point(),
-        qx.suggest_memory_format());
-    using Vec = Vec256<scalar_t>;
-    auto iter = TensorIterator::unary_op(qy, qx);
-    auto zero_point_vec = Vec(scalar_t(zero_point));
-    cpu_kernel_vec(
-        iter,
-        [&](scalar_t value) -> scalar_t {
-          return scalar_t(std::max<underlying_t>(value.val_, zero_point));
-        },
-        [&](Vec value) -> Vec { return value.relu(zero_point_vec); });
-  });
+  qrelu_stub(qx.device().type(), qx, qy);
   return qy;
 }
 Tensor& quantized_relu_(Tensor& qx) {
@@ -50,33 +38,10 @@ Tensor& quantized_relu_(Tensor& qx) {
 namespace {
 Tensor quantized_relu6(const Tensor& qx) {
   Tensor qy;
-  const auto zero_point = qx.q_zero_point();
-  AT_DISPATCH_QINT_TYPES(qx.scalar_type(), "qrelu", [&]() {
-    qy = at::_empty_affine_quantized(
-        qx.sizes(),
-        at::device(kCPU).dtype(SCALAR_TYPE),
-        qx.q_scale(),
-        qx.q_zero_point(),
-        qx.suggest_memory_format());
-    using Vec = Vec256<scalar_t>;
-    auto iter = TensorIterator::unary_op(qy, qx);
-    scalar_t six = at::quantize_val<scalar_t>(qx.q_scale(), qx.q_zero_point(),
-                                              6.0);
-    auto zero_point_vec = Vec(scalar_t(zero_point));
-    auto six_vec = Vec(six);
-    cpu_kernel_vec(
-        iter,
-        [&](scalar_t value) -> scalar_t {
-          underlying_t relu_val =
-              std::max<underlying_t>(value.val_, zero_point);
-          return scalar_t(std::min<underlying_t>(relu_val, six.val_));
-        },
-        [&](Vec val) { return val.relu6(zero_point_vec, six_vec); });
-  });
+  qrelu6_stub(qx.device().type(), qx, qy);
   return qy;
 }
 
-
 class QRelu6 final : public c10::OperatorKernel {
  public:
   Tensor operator()(Tensor qx) {

aten/src/ATen/native/quantized/cpu/quantized_ops.h

@@ -0,0 +1,18 @@
+#include <ATen/ATen.h>
+#include <ATen/native/DispatchStub.h>
+#include <ATen/native/TensorIterator.h>
+
+namespace at {
+namespace native {
+
+using qrelu_fn = void (*)(const at::Tensor& /*qx*/, at::Tensor& /*qy*/);
+using qadd_fn =
+    void (*)(Tensor& /*out*/, const Tensor& /*self*/, const Tensor& /*other*/);
+
+DECLARE_DISPATCH(qrelu_fn, qrelu_stub);
+DECLARE_DISPATCH(qrelu_fn, qrelu6_stub);
+DECLARE_DISPATCH(qadd_fn, qadd_stub);
+DECLARE_DISPATCH(qadd_fn, qadd_relu_stub);
+
+} // namespace native
+} // namespace at

cmake/Codegen.cmake

@@ -84,7 +84,7 @@ if (INTERN_BUILD_ATEN_OPS)
     SET_SOURCE_FILES_PROPERTIES(${CMAKE_CURRENT_LIST_DIR}/../aten/src/TH/THAllocator.cpp PROPERTIES COMPILE_FLAGS "-fno-openmp")
   ENDIF()
 
-  FILE(GLOB cpu_kernel_cpp_in "${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen/native/cpu/*.cpp")
+  FILE(GLOB cpu_kernel_cpp_in "${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen/native/cpu/*.cpp" "${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen/native/quantized/cpu/kernels/*.cpp")
 
   LIST(APPEND CPU_CAPABILITY_NAMES "DEFAULT")
   LIST(APPEND CPU_CAPABILITY_FLAGS "${OPT_FLAG}")

test/test_quantized.py

@@ -80,13 +80,10 @@ def pool_output_shape(input_size, kernel_size, padding, stride,
 class TestQuantizedOps(TestCase):
 
     """Tests the correctness of the quantized::relu op."""
-    @given(qparams=hu.qparams())
-    def test_qrelu(self, qparams):
-        X = np.array([[-3, -2, 1, 2],
-                      [0, 0, 0, 0],
-                      [-5, -4, -3, -2],
-                      [1, 2, 3, 4]], dtype=np.float32)
-        scale, zero_point, torch_type = qparams
+    @given(X=hu.tensor(shapes=hu.array_shapes(1, 5, 1, 5),
+                       qparams=hu.qparams()))
+    def test_qrelu(self, X):
+        X, (scale, zero_point, torch_type) = X
 
         Y = X.copy()
         Y[Y < 0] = 0