[reland] Unify Quantization APIs for add, pool and relu

aten/src/ATen/Context.h

@@ -130,8 +130,6 @@ class CAFFE2_API Context {
   at::QEngine quantized_engine =
 #ifdef USE_FBGEMM
       at::kFBGEMM;
-#elif defined(USE_PYTORCH_QNNPACK)
-      at::kQNNPACK;
 #else
       at::kNoQEngine;
 #endif

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

@@ -5,6 +5,9 @@
 #include <ATen/native/cpu/Loops.h>
 #include <ATen/quantized/Quantizer.h>
 #include <ATen/native/quantized/cpu/quantized_ops.h>
+#include <ATen/native/quantized/cpu/init_qnnpack.h>
+#include <ATen/native/quantized/cpu/qnnpack_utils.h>
+#include <caffe2/utils/threadpool/ThreadPoolMobile.h>
 
 #include <algorithm>
 
@@ -67,9 +70,82 @@ Tensor _add_scalar_out(Tensor& out, const Tensor& self, Scalar other) {
 
 template <bool ReLUFused = false>
 class QAdd final : public c10::OperatorKernel {
+#ifdef USE_PYTORCH_QNNPACK
+Tensor qnnpack_add(Tensor qa, Tensor qb, double scale, int64_t zero_point) {
+  TORCH_CHECK(qa.ndimension() > 0, "qnnpack_add(): Got empty input tensor.");
+  Tensor qa_contig = qa.contiguous();
+  Tensor qb_contig = qb.contiguous();
+
+  const auto a_zero_point = qa_contig.q_zero_point();
+  const auto b_zero_point = qb_contig.q_zero_point();
+  const auto a_scale = qa_contig.q_scale();
+  const auto b_scale = qb_contig.q_scale();
+
+  Tensor qy = at::_empty_affine_quantized(
+      qa_contig.sizes(), at::device(kCPU).dtype(kQUInt8), scale, zero_point);
+
+  if (qa_contig.size(0) == 0) {
+    return qy;
+  }
+
+  initQNNPACK();
+
+  pytorch_qnnp_operator_t qnnpack_operator{nullptr};
+
+  size_t num_elems = qa_contig.numel() / qa_contig.size(0);
+
+  const pytorch_qnnp_status createStatus = pytorch_qnnp_create_add_nc_q8(
+      num_elems /* input size */,
+      a_zero_point /* a zero_point */,
+      a_scale /* a scale */,
+      b_zero_point /* b zero_point */,
+      b_scale /* b scale */,
+      static_cast<uint8_t>(zero_point) /* sum zero_point */,
+      scale /* sum scale */,
+      std::numeric_limits<uint8_t>::min() /* output min */,
+      std::numeric_limits<uint8_t>::max() /* output max */,
+      0 /* flags */,
+      &qnnpack_operator);
+
+  TORCH_INTERNAL_ASSERT(
+      createStatus == pytorch_qnnp_status_success,
+      "failed to create QNNPACK Add operator");
+
+  std::unique_ptr<pytorch_qnnp_operator, QnnpackOperatorDeleter>
+      qnnpack_uniq_ptr(qnnpack_operator);
+
+  const pytorch_qnnp_status setupStatus = pytorch_qnnp_setup_add_nc_q8(
+      qnnpack_operator /* add op */,
+      qa_contig.size(0) /* batch size */,
+      (uint8_t*)qa_contig.data_ptr<c10::quint8>() /* a data */,
+      num_elems /* A stride */,
+      (uint8_t*)qb_contig.data_ptr<c10::quint8>() /* b data */,
+      num_elems /* B stride */,
+      (uint8_t*)qy.data_ptr<c10::quint8>() /* output data */,
+      num_elems /* sum stride */);
+  TORCH_INTERNAL_ASSERT(
+      setupStatus == pytorch_qnnp_status_success,
+      "failed to setup QNNPACK Add operator");
+
+  pthreadpool_t threadpool = caffe2::mobile_threadpool();
+  const pytorch_qnnp_status runStatus =
+      pytorch_qnnp_run_operator(qnnpack_operator, threadpool);
+
+  TORCH_INTERNAL_ASSERT(
+      runStatus == pytorch_qnnp_status_success,
+      "failed to run QNNPACK Add operator");
+
+  return qy;
+}
+#endif
  public:
   Tensor operator()(Tensor qa, Tensor qb, double scale, int64_t zero_point) {
     check_inputs(qa, qb);
+    #ifdef USE_PYTORCH_QNNPACK
+    if (at::globalContext().qEngine() == at::QEngine::QNNPACK) {
+      return qnnpack_add(qa, qb, scale, zero_point);
+    }
+    #endif
     auto qc = at::_empty_affine_quantized(
         qa.sizes(),
         at::device(kCPU).dtype(qa.scalar_type()),