[oneDNN] Support optional scale/shift from oneDNN backend (#58488)

LLee233 · web-flow · commit 463dfdcd312c · 2023-11-02T16:46:00.000+08:00
diff --git a/paddle/phi/backends/onednn/onednn_reuse.h b/paddle/phi/backends/onednn/onednn_reuse.h
@@ -1327,6 +1327,8 @@ class BatchNormOneDNNHandler
                          Place cpu_place,
                          const DenseTensor* x,
                          const float epsilon,
+                         const bool use_scale,
+                         const bool use_bias,
                          const bool fuse_with_relu,
                          const bool global_stats,
                          const bool test_mode)
@@ -1335,8 +1337,9 @@ class BatchNormOneDNNHandler
                                 dnnl::batch_normalization_backward>(engine,
                                                                     cpu_place) {
     // Flags are added by bitwise OR operation
-    auto flags = dnnl::normalization_flags::use_scale |
-                 dnnl::normalization_flags::use_shift;
+    auto flags = dnnl::normalization_flags::none;
+    if (use_scale) flags |= dnnl::normalization_flags::use_scale;
+    if (use_bias) flags |= dnnl::normalization_flags::use_shift;
     if (global_stats) flags |= dnnl::normalization_flags::use_global_stats;
     if (fuse_with_relu && test_mode)
       flags |= dnnl::normalization_flags::fuse_norm_relu;
@@ -1354,39 +1357,31 @@ class BatchNormOneDNNHandler
                          Place cpu_place,
                          const float epsilon,
                          const DenseTensor* in_x,
-                         const DenseTensor* scale,
+                         const bool use_scale,
+                         const bool use_bias,
                          const DenseTensor* out_grad)
       : OneDNNHandlerNoCachingT<T,
                                 dnnl::batch_normalization_forward,
                                 dnnl::batch_normalization_backward>(engine,
                                                                     cpu_place) {
-    auto scale_tz = vectorize<int64_t>(scale->dims());
-    PADDLE_ENFORCE_EQ(
-        scale_tz.size(),
-        1,
-        errors::InvalidArgument(
-            "Dims of scale tensor must be 1, but received scale's size is %d",
-            scale_tz.size()));
+    auto flags = dnnl::normalization_flags::none;
+    if (use_scale) flags |= dnnl::normalization_flags::use_scale;
+    if (use_bias) flags |= dnnl::normalization_flags::use_shift;
 
-    this->AcquireForwardPrimitiveDescriptor(
-        dnnl::prop_kind::forward_training,
-        in_x->mem_desc(),
-        in_x->mem_desc(),
-        epsilon,
-        dnnl::normalization_flags::use_scale |
-            dnnl::normalization_flags::use_shift);
-    this->AcquireBackwardPrimitiveDescriptor(
-        dnnl::prop_kind::backward,
-        out_grad->mem_desc(),
-        out_grad->mem_desc(),
-        in_x->mem_desc(),
-        epsilon,
-        dnnl::normalization_flags::use_scale |
-            dnnl::normalization_flags::use_shift);
+    this->AcquireForwardPrimitiveDescriptor(dnnl::prop_kind::forward_training,
+                                            in_x->mem_desc(),
+                                            in_x->mem_desc(),
+                                            epsilon,
+                                            flags);
+    this->AcquireBackwardPrimitiveDescriptor(dnnl::prop_kind::backward,
+                                             out_grad->mem_desc(),
+                                             out_grad->mem_desc(),
+                                             in_x->mem_desc(),
+                                             epsilon,
+                                             flags);
   }
 
-  std::tuple<std::shared_ptr<dnnl::memory>, std::shared_ptr<dnnl::memory>>
-  AcquireScaleShiftMemory(const DenseTensor* scale, const DenseTensor* shift) {
+  std::shared_ptr<dnnl::memory> AcquireScaleMemory(const DenseTensor* scale) {
     auto scale_tz = vectorize(scale->dims());
     PADDLE_ENFORCE_EQ(
         scale_tz.size(),
@@ -1397,20 +1392,37 @@ class BatchNormOneDNNHandler
 
     auto scale_memory = this->AcquireMemoryFromPrimitive(
         this->fwd_pd_->weights_desc(), to_void_cast<T>(scale->data<T>()));
+
+    return scale_memory;
+  }
+
+  std::shared_ptr<dnnl::memory> AcquireShiftMemory(const DenseTensor* shift) {
+    auto shift_tz = vectorize(shift->dims());
+    PADDLE_ENFORCE_EQ(
+        shift_tz.size(),
+        1,
+        errors::InvalidArgument(
+            "Dims of bias tensor must be 1, but received bias's size is %d",
+            shift_tz.size()));
+
     auto shift_memory = this->AcquireMemoryFromPrimitive(
         this->fwd_pd_->weights_desc(), to_void_cast<T>(shift->data<T>()));
 
-    return std::make_tuple(scale_memory, shift_memory);
+    return shift_memory;
   }
 
-  std::tuple<std::shared_ptr<dnnl::memory>, std::shared_ptr<dnnl::memory>>
-  AcquireDiffScaleShiftMemory(T* diff_scale_data, T* diff_shift_data) {
+  std::shared_ptr<dnnl::memory> AcquireDiffScaleMemory(T* diff_scale_data) {
     auto diff_scale_memory = this->AcquireMemoryFromPrimitive(
         this->bwd_pd_->diff_weights_desc(), diff_scale_data);
+
+    return diff_scale_memory;
+  }
+
+  std::shared_ptr<dnnl::memory> AcquireDiffShiftMemory(T* diff_shift_data) {
     auto diff_shift_memory = this->AcquireMemoryFromPrimitive(
         this->bwd_pd_->diff_weights_desc(), diff_shift_data);
 
-    return std::make_tuple(diff_scale_memory, diff_shift_memory);
+    return diff_shift_memory;
   }
 
   std::shared_ptr<dnnl::memory> AcquireMeanMemory(const DenseTensor* mean) {
diff --git a/paddle/phi/kernels/onednn/batch_norm_grad_kernel.cc b/paddle/phi/kernels/onednn/batch_norm_grad_kernel.cc
@@ -65,8 +65,37 @@ void BatchNormGradFunctor(const Context& dev_ctx,
                           DenseTensor* bias_grad) {
   auto Scale = scale.get_ptr();
   auto Bias = bias.get_ptr();
-  funcs::BatchNormOneDNNHandler<T> handler(
-      dev_ctx.GetEngine(), dev_ctx.GetPlace(), epsilon, &x, Scale, &y_grad);
+  const bool use_scale = scale ? true : false;
+  const bool use_bias = bias ? true : false;
+
+  std::vector<int64_t> scale_tz;
+  std::vector<int64_t> bias_tz;
+  if (use_scale) {
+    scale_tz = vectorize<int64_t>(Scale->dims());
+    PADDLE_ENFORCE_EQ(
+        scale_tz.size(),
+        1,
+        errors::InvalidArgument(
+            "Dims of scale tensor must be 1, but received scale's size is %d",
+            scale_tz.size()));
+  }
+  if (use_bias) {
+    bias_tz = vectorize<int64_t>(Bias->dims());
+    PADDLE_ENFORCE_EQ(
+        bias_tz.size(),
+        1,
+        errors::InvalidArgument(
+            "Dims of bias tensor must be 1, but received bias's size is %d",
+            bias_tz.size()));
+  }
+
+  funcs::BatchNormOneDNNHandler<T> handler(dev_ctx.GetEngine(),
+                                           dev_ctx.GetPlace(),
+                                           epsilon,
+                                           &x,
+                                           use_scale,
+                                           use_bias,
+                                           &y_grad);
 
   T* diff_scale_data = dev_ctx.template Alloc<T>(scale_grad);
   T* diff_shift_data = dev_ctx.template Alloc<T>(bias_grad);
@@ -75,24 +104,29 @@ void BatchNormGradFunctor(const Context& dev_ctx,
   auto mean_memory = handler.AcquireMeanMemory(&saved_mean);
   auto variance_memory = handler.AcquireVarianceMemory(&saved_variance);
   auto diff_dst_memory = handler.AcquireDiffDstMemory(&y_grad);
-  auto scaleshift_mems = handler.AcquireScaleShiftMemory(Scale, Bias);
   auto diff_src_memory = handler.AcquireDiffSrcMemory(x_grad);
-  auto diff_scaleshift_mems =
-      handler.AcquireDiffScaleShiftMemory(diff_scale_data, diff_shift_data);
 
   auto batch_norm_bwd_p = handler.AcquireBackwardPrimitive();
 
+  std::shared_ptr<dnnl::memory> scale_memory(nullptr);
+  std::shared_ptr<dnnl::memory> diff_scale_memory(nullptr);
+  std::shared_ptr<dnnl::memory> diff_shift_memory(nullptr);
+  if (scale) {
+    scale_memory = handler.AcquireScaleMemory(Scale);
+    diff_scale_memory = handler.AcquireDiffScaleMemory(diff_scale_data);
+  }
+  if (bias) diff_shift_memory = handler.AcquireDiffShiftMemory(diff_shift_data);
+
   auto& astream = OneDNNContext::tls().get_stream();
-  batch_norm_bwd_p->execute(
-      astream,
-      {{DNNL_ARG_SRC, *src_memory},
-       {DNNL_ARG_MEAN, *mean_memory},
-       {DNNL_ARG_VARIANCE, *variance_memory},
-       {DNNL_ARG_DIFF_DST, *diff_dst_memory},
-       {DNNL_ARG_SCALE, *(std::get<0>(scaleshift_mems))},
-       {DNNL_ARG_DIFF_SRC, *diff_src_memory},
-       {DNNL_ARG_DIFF_SCALE, *(std::get<0>(diff_scaleshift_mems))},
-       {DNNL_ARG_DIFF_SHIFT, *(std::get<1>(diff_scaleshift_mems))}});
+  batch_norm_bwd_p->execute(astream,
+                            {{DNNL_ARG_SRC, *src_memory},
+                             {DNNL_ARG_MEAN, *mean_memory},
+                             {DNNL_ARG_VARIANCE, *variance_memory},
+                             {DNNL_ARG_DIFF_DST, *diff_dst_memory},
+                             {DNNL_ARG_SCALE, *scale_memory},
+                             {DNNL_ARG_DIFF_SRC, *diff_src_memory},
+                             {DNNL_ARG_DIFF_SCALE, *diff_scale_memory},
+                             {DNNL_ARG_DIFF_SHIFT, *diff_shift_memory}});
   astream.wait();
 
   // set memory descriptor of out tensor
diff --git a/paddle/phi/kernels/onednn/batch_norm_kernel.cc b/paddle/phi/kernels/onednn/batch_norm_kernel.cc
@@ -49,19 +49,20 @@ void BatchNormKernel(const Context &dev_ctx,
       dev_ctx.HasDnnAttr("fuse_with_relu")
           ? PADDLE_GET_CONST(bool, dev_ctx.GetDnnAttr("fuse_with_relu"))
           : false;
+  const bool use_scale = scale ? true : false;
+  const bool use_bias = bias ? true : false;
 
-  auto Scale = scale.get_ptr();
-  auto Bias = bias.get_ptr();
   funcs::BatchNormOneDNNHandler<T> handler(dev_ctx.GetEngine(),
                                            dev_ctx.GetPlace(),
                                            &x,
                                            epsilon,
+                                           use_scale,
+                                           use_bias,
                                            fuse_with_relu,
                                            global_stats,
                                            test_mode);
 
   auto src_memory = handler.AcquireSrcMemory(&x);
-  auto scaleshift_mems = handler.AcquireScaleShiftMemory(Scale, Bias);
   auto dst_memory = handler.AcquireDstMemory(y);
   auto batch_norm_p = handler.AcquireForwardPrimitive();
 
@@ -79,18 +80,25 @@ void BatchNormKernel(const Context &dev_ctx,
 
   y->set_mem_desc(dst_memory->get_desc());
 
+  std::shared_ptr<dnnl::memory> scale_memory(nullptr);
+  std::shared_ptr<dnnl::memory> shift_memory(nullptr);
+  auto Scale = scale.get_ptr();
+  auto Bias = bias.get_ptr();
+  if (scale) scale_memory = handler.AcquireScaleMemory(Scale);
+  if (bias) shift_memory = handler.AcquireShiftMemory(Bias);
+
   auto &astream = OneDNNContext::tls().get_stream();
   batch_norm_p->execute(astream,
                         {{DNNL_ARG_SRC, *src_memory},
-                         {DNNL_ARG_SCALE, *(std::get<0>(scaleshift_mems))},
-                         {DNNL_ARG_SHIFT, *(std::get<1>(scaleshift_mems))},
+                         {DNNL_ARG_SCALE, *scale_memory},
+                         {DNNL_ARG_SHIFT, *shift_memory},
                          {DNNL_ARG_MEAN, *mean_memory},
                          {DNNL_ARG_VARIANCE, *variance_memory},
                          {DNNL_ARG_DST, *dst_memory}});
   astream.wait();
 
   if (!global_stats) {
-    const unsigned int C = phi::vectorize(Scale->dims())[0];
+    const unsigned int C = phi::vectorize(mean.dims())[0];
 
     // mkldnn only compute stats for current batch
     // so we need compute momentum stats via Eigen lib
