Rc4/separate memory management for quantizer (#2580)

zhuyuhua-v · ZhiweiYan-96 · fengyuan14 · web-flow · commit 9a407189eaa1 · 2023-04-21T15:28:28.000+08:00
* Refactor q reorder design (#2471)
* Refactor q reorder design
* Initialize args mark in constructor
* Remove exclusive relation on sc&amp;zp setting between src&amp;dst
* Change quantized_reorder api
* Use cached sc&amp;zp in quant/dequant

---------

Co-authored-by: Zhiwei &lt;zhiwei.yan@intel.com&gt;
Co-authored-by: Feng Yuan &lt;feng1.yuan@intel.com&gt;
Co-authored-by: xiaolil1 &lt;xiaoli.liu@intel.com&gt;
diff --git a/csrc/gpu/aten/operators/ReQuantization.cpp b/csrc/gpu/aten/operators/ReQuantization.cpp
@@ -25,16 +25,21 @@ Tensor requantize(
   auto reorder_attr = xpu::oneDNN::ReorderAttr();
   int mask = 0;
   auto scale_in = src.is_quantized() ? static_cast<float>(src.q_scale()) : 1.f;
-  auto requant_scale = static_cast<float>(1.f / (scale_out / scale_in));
+  auto requant_scale = static_cast<float>((scale_out / scale_in));
 
-  Tensor dnn_scale =
-      at::ones(1, at::dtype(at::kFloat).device(at::kXPU)) * requant_scale;
+  Tensor dnn_scale = at::empty({1}, at::dtype(at::kFloat).device(at::kXPU))
+                         .fill_(requant_scale);
   // TODO: Remove workaround for dnnl symmetric quantization
-  Tensor dnn_zero_point =
-      at::ones(1, at::dtype(at::kInt).device(at::kXPU)) * zero_point_out;
+  Tensor dnn_zero_point = at::zeros({1}, at::dtype(at::kInt).device(at::kXPU));
   reorder_attr.set_dst_sc_and_zp_mask(mask);
   xpu::oneDNN::quantized_reorder(
-      src, dst_, dnn_scale, dnn_zero_point, reorder_attr);
+      src,
+      dst_,
+      /*src_scale=*/Tensor(),
+      /*src_zero_point=*/Tensor(),
+      dnn_scale,
+      dnn_zero_point,
+      reorder_attr);
 
   return dst_;
 }
diff --git a/csrc/gpu/aten/quantized/DeQuantization.cpp b/csrc/gpu/aten/quantized/DeQuantization.cpp
@@ -14,6 +14,7 @@ using namespace dnnl;
 using namespace at::native;
 using namespace xpu::dpcpp;
 using namespace xpu::oneDNN;
+using namespace at::AtenIpexTypeQuantizedXPU;
 
 namespace at {
 namespace AtenIpexTypeXPU {
@@ -25,25 +26,38 @@ Tensor dequantize_tensor_per_tensor_affine(
     int64_t zero_point) {
   ReorderAttr rattr = ReorderAttr();
   int mask = 0;
-  auto q_ctx = DPCPPTensorContext::get_tensor_ctx(qtensor);
-  // TODO: Remove workaround for dnnl symmetric quantization
-  float true_scale = ((q_ctx.is_plain() ? get_onednn_dtype(qtensor)
-                                        : q_ctx.meta().get_data_type()) ==
-                          memory::data_type::u8 &&
-                      qtensor.q_zero_point() == 128)
-      ? static_cast<float>(scale / 2)
-      : static_cast<float>(scale);
   rattr.set_src_sc_and_zp_mask(mask);
 
-  // See [Note: Scale setting for reorder]
-  Tensor dnn_scale =
-      at::ones(1, at::dtype(at::kFloat).device(at::kXPU)) * true_scale;
-  // TODO: Remove workaround for dnnl symmetric quantization
-  Tensor dnn_zero_point = at::zeros(1, at::dtype(at::kInt).device(at::kXPU));
-
   Tensor rtensor_ = at::empty(qtensor.sizes(), rtensor.options());
-  xpu::oneDNN::quantized_reorder(
-      qtensor, rtensor_, dnn_scale, dnn_zero_point, rattr);
+  if (is_opaque_u8(qtensor)) {
+    Tensor dnn_scale =
+        at::empty({1}, at::dtype(at::kFloat).device(at::kXPU)).fill_(scale);
+    Tensor dnn_zero_point =
+        at::zeros({1}, at::dtype(at::kInt).device(at::kXPU));
+
+    // See [Note: Scale setting for reorder]
+    xpu::oneDNN::quantized_reorder(
+        qtensor,
+        rtensor_,
+        dnn_scale,
+        dnn_zero_point,
+        /*dst_scale=*/Tensor(),
+        /*dst_zero_point=*/Tensor(),
+        rattr);
+  } else {
+    // See [Note: Scale setting for reorder]
+    xpu::oneDNN::quantized_reorder(
+        qtensor,
+        rtensor_,
+        q_scale_ptr(qtensor),
+        q_zero_point_ptr(qtensor),
+        /*dst_scale=*/nullptr,
+        /*dst_zero_point=*/nullptr,
+        {1},
+        {1},
+        rattr);
+  }
+
   return rtensor_;
 }
 
@@ -91,7 +105,13 @@ Tensor dequantize_tensor_per_channel_affine(
 
   Tensor rtensor_ = empty_opaque_tensor(r_md, rtensor.options(), c10::nullopt);
   xpu::oneDNN::quantized_reorder(
-      qtensor, rtensor_, dnn_scale, dnn_zero_point, rattr);
+      qtensor,
+      rtensor_,
+      dnn_scale,
+      dnn_zero_point,
+      /*dst_scale=*/Tensor(),
+      /*dst_zero_point=*/Tensor(),
+      rattr);
 
   return rtensor_;
 }
diff --git a/csrc/gpu/aten/quantized/QTensor.cpp b/csrc/gpu/aten/quantized/QTensor.cpp
@@ -122,20 +122,19 @@ int64_t q_per_channel_axis(const Tensor& self) {
   return at::native::q_per_channel_axis(self);
 }
 
-Tensor q_scale_tensor(const Tensor& self) {
+float* q_scale_ptr(const Tensor& self) {
   auto quantizer = get_qtensorimpl(self)->quantizer();
   TORCH_CHECK(quantizer->qscheme() == kPerTensorAffine);
   return static_cast<DPCPPPerTensorAffineQuantizer*>(quantizer.get())
-      ->scale_tensor();
+      ->scale_ptr();
 }
 
-Tensor q_zero_point_tensor(const Tensor& self) {
+int32_t* q_zero_point_ptr(const Tensor& self) {
   auto quantizer = get_qtensorimpl(self)->quantizer();
   TORCH_CHECK(quantizer->qscheme() == kPerTensorAffine);
   return static_cast<DPCPPPerTensorAffineQuantizer*>(quantizer.get())
-      ->zero_point_tensor();
+      ->zero_point_ptr();
 }
-
 Tensor& set_(
     Tensor& self,
     Storage storage,
diff --git a/csrc/gpu/aten/quantized/QTensor.h b/csrc/gpu/aten/quantized/QTensor.h
@@ -36,9 +36,9 @@ Tensor q_per_channel_zero_points(const Tensor& self);
 
 int64_t q_per_channel_axis(const Tensor& self);
 
-Tensor q_scale_tensor(const Tensor& self);
+float* q_scale_ptr(const Tensor& self);
 
-Tensor q_zero_point_tensor(const Tensor& self);
+int32_t* q_zero_point_ptr(const Tensor& self);
 
 Tensor& set_(
     Tensor& self,
diff --git a/csrc/gpu/aten/quantized/Quantization.cpp b/csrc/gpu/aten/quantized/Quantization.cpp
@@ -77,7 +77,13 @@ Tensor quantize_tensor_per_channel_affine(
   Tensor dnn_zero_point =
       at::zeros_like(zero_points, dtype(at::kInt).device(at::kXPU));
   xpu::oneDNN::quantized_reorder(
-      rtensor, qtensor, dnn_scale, dnn_zero_point, rattr);
+      rtensor,
+      qtensor,
+      /*src_scale=*/Tensor(),
+      /*src_zero_point=*/Tensor(),
+      dnn_scale,
+      dnn_zero_point,
+      rattr);
 
   return qtensor;
 }
@@ -134,15 +140,27 @@ Tensor quantize_tensor_per_tensor_affine(
         AtenIpexTypeXPU::empty_opaque_qtensor(q_md, c10::nullopt, quantizer);
 
     xpu::oneDNN::quantized_reorder(
-        rtensor, qtensor_opt, dnn_scale, dnn_zero_point, rattr);
+        rtensor,
+        qtensor_opt,
+        /*src_scale=*/Tensor(),
+        /*src_zero_point=*/Tensor(),
+        dnn_scale,
+        dnn_zero_point,
+        rattr);
     auto q_opt_ctx =
         at::AtenIpexTypeXPU::DPCPPTensorContext::release_tensor_ctx(
             qtensor_opt);
     at::AtenIpexTypeXPU::DPCPPTensorContext::set_tensor_ctx(
         qtensor, std::move(q_opt_ctx));
   } else {
     xpu::oneDNN::quantized_reorder(
-        rtensor, qtensor, dnn_scale, dnn_zero_point, rattr);
+        rtensor,
+        qtensor,
+        /*src_scale=*/Tensor(),
+        /*srd_zero_point=*/Tensor(),
+        dnn_scale,
+        dnn_zero_point,
+        rattr);
   }
 
   return qtensor;
diff --git a/csrc/gpu/aten/quantized/Quantizer.h b/csrc/gpu/aten/quantized/Quantizer.h
@@ -4,12 +4,50 @@
 #include <quantized/DeQuantization.h>
 #include <quantized/QTensor.h>
 #include <quantized/Quantization.h>
+#include <runtime/Utils.h>
 #include <utils/LRUCache.h>
 
 namespace at {
 namespace AtenIpexTypeQuantizedXPU {
 
+using namespace xpu::dpcpp;
+
+template <typename scale_t_, typename zp_t_>
+class XPUQuantizerBase {
+ public:
+  using scale_t = scale_t_;
+  using zp_t = zp_t_;
+  using scale_ptr_t = std::shared_ptr<scale_t>;
+  using zp_ptr_t = std::shared_ptr<zp_t>;
+
+ public:
+  XPUQuantizerBase() = default;
+
+  XPUQuantizerBase(size_t size, sycl::queue& q)
+      : scale_ptr_(
+            sycl::malloc_device<scale_t>(size * sizeof(scale_t), q),
+            [=](scale_t* ptr) { sycl::free(ptr, q); }),
+        zp_ptr_(
+            sycl::malloc_device<zp_t>(size * sizeof(zp_t), q),
+            [=](zp_t* ptr) { sycl::free(ptr, q); }) {}
+  scale_t* scale_ptr() {
+    return scale_ptr_.get();
+  }
+
+  zp_t* zero_point_ptr() {
+    return zp_ptr_.get();
+  }
+
+ private:
+  scale_ptr_t scale_ptr_;
+  zp_ptr_t zp_ptr_;
+};
+
 struct DPCPPPerTensorAffineQuantizer : public AffineQuantizer {
+  using QuantizerBaseType = XPUQuantizerBase<float, int32_t>;
+  using scale_t = QuantizerBaseType::scale_t;
+  using zp_t = QuantizerBaseType::zp_t;
+
   explicit DPCPPPerTensorAffineQuantizer(
       ScalarType scalar_type,
       double scale,
@@ -22,17 +60,21 @@ struct DPCPPPerTensorAffineQuantizer : public AffineQuantizer {
     }
     // TODO: Modify this line after asymmetric enabled
     xpu::dpcpp::create_key(key_sc_zp, dnn_scale, 0);
-    bool key_found = xpu::dpcpp::find_key<std::pair<Tensor, Tensor>>(key_sc_zp);
+    bool key_found = xpu::dpcpp::find_key<QuantizerBaseType>(key_sc_zp);
     if (key_found) {
-      std::tie(scale_tensor_, zero_point_tensor_) =
-          xpu::dpcpp::fetch_m<std::pair<Tensor, Tensor>>(key_sc_zp);
+      base_ = xpu::dpcpp::fetch_m<QuantizerBaseType>(key_sc_zp);
     } else {
-      scale_tensor_ = at::empty({1}, at::dtype(kFloat).device(at::kXPU))
-                          .fill_(static_cast<float>(dnn_scale));
-      // TODO: Modify this line after asymmetric enabled
-      zero_point_tensor_ = at::zeros({1}, at::dtype(kInt).device(at::kXPU));
-      xpu::dpcpp::fetch_or_create_m<std::pair<Tensor, Tensor>>(
-          key_sc_zp, scale_tensor_, zero_point_tensor_);
+      base_ = QuantizerBaseType(1, dpcppGetCurrentQueue());
+
+      scale_t* sc_ptr = base_.scale_ptr();
+      scale_t _scale = (scale_t)dnn_scale;
+      dpcppGetCurrentQueue().single_task([=]() { sc_ptr[0] = _scale; });
+
+      zp_t* zp_ptr = base_.zero_point_ptr();
+      zp_t _zp = (zp_t)0;
+      dpcppGetCurrentQueue().single_task([=]() { zp_ptr[0] = _zp; });
+
+      xpu::dpcpp::fetch_or_create_m<QuantizerBaseType>(key_sc_zp, base_);
     }
   }
 
@@ -80,12 +122,12 @@ struct DPCPPPerTensorAffineQuantizer : public AffineQuantizer {
     return zero_point_;
   }
 
-  Tensor scale_tensor() {
-    return scale_tensor_;
+  scale_t* scale_ptr() {
+    return base_.scale_ptr();
   }
 
-  Tensor zero_point_tensor() {
-    return zero_point_tensor_;
+  zp_t* zero_point_ptr() {
+    return base_.zero_point_ptr();
   }
 
   bool equalTo(QuantizerPtr other) const override {
@@ -107,8 +149,7 @@ struct DPCPPPerTensorAffineQuantizer : public AffineQuantizer {
   const double scale_;
   // We use int64_t for consistency with Python
   const int64_t zero_point_;
-  Tensor scale_tensor_;
-  Tensor zero_point_tensor_;
+  QuantizerBaseType base_;
 };
 
 struct DPCPPPerChannelAffineQuantizer : public AffineQuantizer {
diff --git a/csrc/gpu/oneDNN/Matmul.h b/csrc/gpu/oneDNN/Matmul.h
@@ -363,8 +363,7 @@ static inline void matmul(
                   .fill_(m1.q_scale());
       m1_sc_m = dpcpp_onednn_memory(m1_sc_md, engine, m1_sc.data_ptr());
     } else {
-      m1_sc = at::AtenIpexTypeQuantizedXPU::q_scale_tensor(m1);
-      m1_sc_m = dpcpp_onednn_memory(m1_sc_md, engine, m1_sc.data_ptr());
+      m1_sc_m = dpcpp_onednn_memory(m1_sc_md, engine, q_scale_ptr(m1));
     }
     args.insert({DNNL_ARG_ATTR_SCALES | DNNL_ARG_SRC, m1_sc_m});
 
@@ -373,8 +372,7 @@ static inline void matmul(
     if (m2.is_quantized()) {
       memory::desc m2_sc_md =
           memory::desc({1}, memory::data_type::f32, memory::format_tag::x);
-      m2_sc = at::AtenIpexTypeQuantizedXPU::q_scale_tensor(m2);
-      m2_sc_m = dpcpp_onednn_memory(m2_sc_md, engine, m2_sc.data_ptr());
+      m2_sc_m = dpcpp_onednn_memory(m2_sc_md, engine, q_scale_ptr(m2));
       args.insert({DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS, m2_sc_m});
     }
 
@@ -383,8 +381,7 @@ static inline void matmul(
     if (dst.is_quantized()) {
       memory::desc dst_sc_md =
           memory::desc({1}, memory::data_type::f32, memory::format_tag::x);
-      dst_sc = at::AtenIpexTypeQuantizedXPU::q_scale_tensor(dst);
-      dst_sc_m = dpcpp_onednn_memory(dst_sc_md, engine, dst_sc.data_ptr());
+      dst_sc_m = dpcpp_onednn_memory(dst_sc_md, engine, q_scale_ptr(dst));
       args.insert({DNNL_ARG_ATTR_SCALES | DNNL_ARG_DST, dst_sc_m});
     }
 
@@ -421,9 +418,7 @@ static inline void matmul(
     if (is_per_tensor_quantized) {
       memory::desc wgh_sc_md =
           memory::desc({1}, memory::data_type::f32, memory::format_tag::x);
-      Tensor wgh_sc = at::AtenIpexTypeQuantizedXPU::q_scale_tensor(m2);
-      memory wgh_sc_m =
-          dpcpp_onednn_memory(wgh_sc_md, engine, wgh_sc.data_ptr());
+      memory wgh_sc_m = dpcpp_onednn_memory(wgh_sc_md, engine, q_scale_ptr(m2));
       args.insert({DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS, wgh_sc_m});
 
 #ifdef BUILD_PRIOR_SYMM_QUANT
diff --git a/csrc/gpu/oneDNN/QConv.h b/csrc/gpu/oneDNN/QConv.h
@@ -485,8 +485,7 @@ static at::Tensor quantized_convolution(
                  .fill_(static_cast<float>(src.q_scale()));
     src_sc_m = dpcpp_onednn_memory(src_sc_md, engine, src_sc.data_ptr());
   } else {
-    src_sc = at::AtenIpexTypeQuantizedXPU::q_scale_tensor(src);
-    src_sc_m = dpcpp_onednn_memory(src_sc_md, engine, src_sc.data_ptr());
+    src_sc_m = dpcpp_onednn_memory(src_sc_md, engine, q_scale_ptr(src));
   }
   args.insert({DNNL_ARG_ATTR_SCALES | DNNL_ARG_SRC, src_sc_m});
 
@@ -530,10 +529,9 @@ static at::Tensor quantized_convolution(
 #endif
 
   if (wgh.qscheme() == kPerTensorAffine) {
-    Tensor wgh_sc = at::AtenIpexTypeQuantizedXPU::q_scale_tensor(wgh);
     memory::desc wgh_sc_md =
         memory::desc({1}, memory::data_type::f32, memory::format_tag::x);
-    memory wgh_sc_m = dpcpp_onednn_memory(wgh_sc_md, engine, wgh_sc.data_ptr());
+    memory wgh_sc_m = dpcpp_onednn_memory(wgh_sc_md, engine, q_scale_ptr(wgh));
     args.insert({DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS, wgh_sc_m});
 
 #ifdef BUILD_PRIOR_SYMM_QUANT
diff --git a/csrc/gpu/oneDNN/QDeconv.h b/csrc/gpu/oneDNN/QDeconv.h
@@ -367,8 +367,7 @@ static Tensor quantized_deconvolution(
                  .fill_(static_cast<float>(src.q_scale()));
     src_sc_m = dpcpp_onednn_memory(src_sc_md, engine, src_sc.data_ptr());
   } else {
-    src_sc = at::AtenIpexTypeQuantizedXPU::q_scale_tensor(src);
-    src_sc_m = dpcpp_onednn_memory(src_sc_md, engine, src_sc.data_ptr());
+    src_sc_m = dpcpp_onednn_memory(src_sc_md, engine, q_scale_ptr(src));
   }
   args.insert({DNNL_ARG_ATTR_SCALES | DNNL_ARG_SRC, src_sc_m});
 
@@ -387,10 +386,9 @@ static Tensor quantized_deconvolution(
   }
 #endif
 
-  Tensor dst_sc = at::AtenIpexTypeQuantizedXPU::q_scale_tensor(dst);
   memory::desc dst_sc_md =
       memory::desc({1}, memory::data_type::f32, memory::format_tag::x);
-  memory dst_sc_m = dpcpp_onednn_memory(dst_sc_md, engine, dst_sc.data_ptr());
+  memory dst_sc_m = dpcpp_onednn_memory(dst_sc_md, engine, q_scale_ptr(dst));
   args.insert({DNNL_ARG_ATTR_SCALES | DNNL_ARG_DST, dst_sc_m});
 
 #ifdef BUILD_PRIOR_SYMM_QUANT
@@ -407,10 +405,9 @@ static Tensor quantized_deconvolution(
 #endif
 
   if (wgh.qscheme() == kPerTensorAffine) {
-    Tensor wgh_sc = at::AtenIpexTypeQuantizedXPU::q_scale_tensor(wgh);
     memory::desc wgh_sc_md =
         memory::desc({1}, memory::data_type::f32, memory::format_tag::x);
-    memory wgh_sc_m = dpcpp_onednn_memory(wgh_sc_md, engine, wgh_sc.data_ptr());
+    memory wgh_sc_m = dpcpp_onednn_memory(wgh_sc_md, engine, q_scale_ptr(wgh));
     args.insert({DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS, wgh_sc_m});
 
 #ifdef BUILD_PRIOR_SYMM_QUANT
diff --git a/csrc/gpu/oneDNN/Reorder.h b/csrc/gpu/oneDNN/Reorder.h
