From ff2fba3987ef9eac3bf240b92564e87832a331b4 Mon Sep 17 00:00:00 2001
From: crystal <62974595+Zjq9409@users.noreply.github.com>
Date: Sat, 9 Apr 2022 19:33:07 +0800
Subject: [PATCH] modify the block size of the group_norm backward (#41570)

---
 paddle/fluid/operators/group_norm_op.cu | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/paddle/fluid/operators/group_norm_op.cu b/paddle/fluid/operators/group_norm_op.cu
index c93910bde5a..bb8031b0cc4 100644
--- a/paddle/fluid/operators/group_norm_op.cu
+++ b/paddle/fluid/operators/group_norm_op.cu
@@ -605,8 +605,16 @@ class GroupNormGradKernel<platform::CUDADeviceContext, T>
     int flags =
         (scale_data != nullptr) * kHasScale + (bias_data != nullptr) * kHasBias;
     if (data_layout == DataLayout::kNCHW) {
+      const int max_num_threads = 1024;
+      int max_block_size = std::min(imsize, max_num_threads);
+      int block_size_nchw = 1;
+      while (block_size_nchw < max_block_size) {
+        block_size_nchw *= 2;
+      }
+      block_size_nchw = std::max(block_size_nchw, kps::details::kWarpSize);
+      dim3 blocks(block_size_nchw);
       ScalarGetDsDbCUDAKernel<
-          T><<<x_dims[0] * C, block_size, 0, dev_ctx.stream()>>>(
+          T><<<x_dims[0] * C, blocks, 0, dev_ctx.stream()>>>(
           imsize, x_data, dy_data, ds_data, db_data);
 
       if (d_scale || d_bias) {
-- 
GitLab