PaddlePaddle
diff --git a/‎paddle/fluid/inference/tensorrt/plugin/group_norm_op_plugin.cu
Copy file name to clipboard
+84-44Lines changed: 84 additions & 44 deletions b/‎paddle/fluid/inference/tensorrt/plugin/group_norm_op_plugin.cu
Copy file name to clipboard
+84-44Lines changed: 84 additions & 44 deletions
diff --git a/‎paddle/fluid/inference/tensorrt/plugin/group_norm_op_plugin.h
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions b/‎paddle/fluid/inference/tensorrt/plugin/group_norm_op_plugin.h
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions
@@ -74,7 +74,7 @@ static int32_t findMaxDivisor(int32_t n, int32_t maxAllowedDivisor) {
 
 template <int tTHREADS_PER_BLOCK>
 __global__ void groupNormNCHW32SumKernelQDQ(
-    const GroupNormNHWCParams<__half> params) {
+    const GroupNormNDHWCParams<__half> params) {
   // The object in charge of doing the sums for the different blocks.
   typedef cub::BlockScan<GroupSums, tTHREADS_PER_BLOCK> BlockScan;
 
@@ -90,9 +90,9 @@ __global__ void groupNormNCHW32SumKernelQDQ(
   int32_t ci = blockIdx.x * params.cPerBlock + threadIdx.x * 2;
 
   // The first activation loaded by that block.
-  int32_t hwBegin = blockIdx.y * params.hwPerBlock;
+  int32_t dhwBegin = blockIdx.y * params.dhwPerBlock;
   // The last activation loaded by that block.
-  int32_t hwEnd = min(hwBegin + params.hwPerBlock, params.hw);
+  int32_t dhwEnd = min(dhwBegin + params.dhwPerBlock, params.dhw);
 
   // The sums.
   float sum = 0.F;
@@ -102,13 +102,13 @@ __global__ void groupNormNCHW32SumKernelQDQ(
 
   // nchw32 layout
   // batch offset + channel offset
-  int nc_offset = static_cast<int64_t>(ni) * params.hwc +
-                  ci / 32 * params.hw * 32 + ci % 32;
+  int nc_offset = static_cast<int64_t>(ni) * params.dhwc +
+                  ci / 32 * params.dhw * 32 + ci % 32;
 
   // Iterate over the activations to compute the sums.
-  for (int32_t hwi = hwBegin; hwi < hwEnd; ++hwi) {
+  for (int32_t dhwi = dhwBegin; dhwi < dhwEnd; ++dhwi) {
     // The offset.
-    int64_t offset = nc_offset + static_cast<int64_t>(hwi) * 32;
+    int64_t offset = nc_offset + static_cast<int64_t>(dhwi) * 32;
 
     // Fetch two channels per thread.
     __half2 h2(0, 0);
@@ -166,14 +166,14 @@ __global__ void groupNormNCHW32SumKernelQDQ(
   atomicAdd(&params.redBuffer[(2 * ni + 1) * params.groups + gj], sums.y);
 }
 
-void groupNormNCHW32SumQDQ(const GroupNormNHWCParams<__half> &params,
+void groupNormNCHW32SumQDQ(const GroupNormNDHWCParams<__half> &params,
                            cudaStream_t stream) {
   dim3 grid;
 
   // The number of blocks to compute all the channels.
   grid.x = divUp(params.c, params.cPerBlock);
   // The number of blocks to compute all the activations in a given instance.
-  grid.y = divUp(params.hw, params.hwPerBlock);
+  grid.y = divUp(params.dhw, params.dhwPerBlock);
   // The number of instances.
   grid.z = params.n;
 
@@ -198,7 +198,7 @@ void groupNormNCHW32SumQDQ(const GroupNormNHWCParams<__half> &params,
 
 template <int tTHREADS_PER_BLOCK>
 __global__ void groupNormNCHW32ScaleKernelQDQ(
-    const GroupNormNHWCParams<__half> params) {
+    const GroupNormNDHWCParams<__half> params) {
   // The instance in the batch.
   int32_t ni = blockIdx.z;
   // The channel loaded by that thread (2 channels per thread for F16x2).
@@ -226,25 +226,25 @@ __global__ void groupNormNCHW32ScaleKernelQDQ(
   }
 
   // Compute the mean.
-  float mean = sum * params.invHWC;
+  float mean = sum * params.invDHWC;
   // Compute the variance.
-  float var = sumSq * params.invHWC - (mean * mean);
+  float var = sumSq * params.invDHWC - (mean * mean);
   // Compute the inverse of the stddev.
   float invStdDev = rsqrtf(var + params.eps);
 
   // The first activation loaded by that block.
-  int32_t hwBegin = blockIdx.y * params.hwPerBlock;
+  int32_t dhwBegin = blockIdx.y * params.dhwPerBlock;
   // The last activation loaded by that block.
-  int32_t hwEnd = min(hwBegin + params.hwPerBlock, params.hw);
+  int32_t dhwEnd = min(dhwBegin + params.dhwPerBlock, params.dhw);
 
   // nchw32 layout
-  int c_offset = ci / 32 * params.hw * 32 + ci % 32;
+  int c_offset = ci / 32 * params.dhw * 32 + ci % 32;
 
   // Iterate over the activations to compute the sums.
-  for (int32_t hwi = hwBegin; hwi < hwEnd; ++hwi) {
+  for (int32_t dhwi = dhwBegin; dhwi < dhwEnd; ++dhwi) {
     // The src/dst offset.
-    int64_t offset = static_cast<int64_t>(ni) * params.hwc + c_offset +
-                     static_cast<int64_t>(hwi) * 32;
+    int64_t offset = static_cast<int64_t>(ni) * params.dhwc + c_offset +
+                     static_cast<int64_t>(dhwi) * 32;
 
     // Fetch two channels per thread.
     __half2 h2(0, 0);
@@ -290,14 +290,14 @@ __global__ void groupNormNCHW32ScaleKernelQDQ(
   }
 }
 
-void groupNormNCHW32ScaleQDQ(const GroupNormNHWCParams<__half> &params,
+void groupNormNCHW32ScaleQDQ(const GroupNormNDHWCParams<__half> &params,
                              cudaStream_t stream) {
   dim3 grid;
 
   // The number of blocks to compute all the channels.
   grid.x = divUp(params.c, params.cPerBlock);
   // The number of blocks to compute all the activations in a given instance.
-  grid.y = divUp(params.hw, params.hwPerBlock);
+  grid.y = divUp(params.dhw, params.dhwPerBlock);
   // The number of instances.
   grid.z = params.n;
 
@@ -642,7 +642,7 @@ int GroupNormPluginDynamic::enqueue(
                  DataLayout::kNCHW);
     } else if (input_desc[0].format == nvinfer1::PluginFormat::kHWC8) {
       int32_t cPerBlock = 320;
-      int32_t maxBlocksPerHW = 1024;
+      int32_t maxBlocksPerDHW = 1024;
       switch (input_desc[0].dims.d[1]) {
         case 960:
         case 1920:
@@ -661,6 +661,25 @@ int GroupNormPluginDynamic::enqueue(
       if (cPerBlock > input_desc[0].dims.d[1]) {
         cPerBlock = 8;
       }
+      auto d_dim = input_desc[0].dims.nbDims;
+      params_.n = input_desc[0].dims.d[0];
+      if (d_dim == 3) {
+        params_.c = input_desc[0].dims.d[1];
+        params_.d = 1;
+        params_.h = 1;
+        params_.w = input_desc[0].dims.d[2];
+      } else if (d_dim == 4) {
+        params_.c = input_desc[0].dims.d[1];
+        params_.d = 1;
+        params_.h = input_desc[0].dims.d[2];
+        params_.w = input_desc[0].dims.d[3];
+      } else {
+        // d_dim == 5
+        params_.c = input_desc[0].dims.d[1];
+        params_.d = input_desc[0].dims.d[2];
+        params_.h = input_desc[0].dims.d[3];
+        params_.w = input_desc[0].dims.d[4];
+      }
 
       params_.withSilu = with_silu_;
       params_.dst = static_cast<half *>(outputs[0]);
@@ -669,18 +688,19 @@ int GroupNormPluginDynamic::enqueue(
       params_.beta = reinterpret_cast<half *>(bias_gpu_);
       params_.redBuffer = static_cast<float *>(workspace);
       params_.var_data = nullptr;
-      params_.n = input_desc[0].dims.d[0];
-      params_.h = input_desc[0].dims.d[2];
-      params_.w = input_desc[0].dims.d[3];
-      params_.c = input_desc[0].dims.d[1];
+      // params_.n = input_desc[0].dims.d[0];
+      // params_.h = input_desc[0].dims.d[2];
+      // params_.w = input_desc[0].dims.d[3];
+      // params_.c = input_desc[0].dims.d[1];
       params_.groups = groups_;
-      params_.hw = params_.h * params_.w;
-      const int32_t blocksPerHW = findMaxDivisor(params_.hw, maxBlocksPerHW);
-      params_.hwPerBlock = divUp(params_.hw, blocksPerHW);
+      params_.dhw = params_.d * params_.h * params_.w;
+      const int32_t blocksPerDHW = findMaxDivisor(params_.dhw, maxBlocksPerDHW);
+      params_.dhwPerBlock = divUp(params_.dhw, blocksPerDHW);
       params_.cPerBlock = cPerBlock;
       params_.cPerGroup = params_.c / params_.groups;
-      params_.hwc = params_.hw * params_.c;
-      params_.invHWC = 1.F / static_cast<float>(params_.hw * params_.cPerGroup);
+      params_.dhwc = params_.dhw * params_.c;
+      params_.invDHWC =
+          1.F / static_cast<float>(params_.dhw * params_.cPerGroup);
       params_.groupsPerBlock = cPerBlock / params_.cPerGroup;
       params_.eps = eps_;
       params_.var_data = nullptr;
@@ -690,10 +710,10 @@ int GroupNormPluginDynamic::enqueue(
                       2 * sizeof(float) * params_.n * groups_,
                       stream);
 
-      phi::groupNormNHWCSum<half> nhwc_sum;
-      nhwc_sum(&params_, stream);
-      phi::groupNormNHWCScale<half> nhwc_scale;
-      nhwc_scale(params_, stream);
+      phi::groupNormNDHWCSum<half> ndhwc_sum;
+      ndhwc_sum(&params_, stream);
+      phi::groupNormNDHWCScale<half> ndhwc_scale;
+      ndhwc_scale(params_, stream);
     } else {
       PADDLE_THROW(platform::errors::Fatal(
           "The Groupnorm TRT Plugin's only support nchw or nhwc8 input"));
@@ -704,7 +724,7 @@ int GroupNormPluginDynamic::enqueue(
 
     if (input_desc[0].format == nvinfer1::PluginFormat::kCHW32) {
       int32_t cPerBlock = 320;
-      int32_t maxBlocksPerHW = 1024;
+      int32_t maxBlocksPerDHW = 1024;
       switch (input_desc[0].dims.d[1]) {
         case 960:
         case 1920:
@@ -723,25 +743,45 @@ int GroupNormPluginDynamic::enqueue(
       if (cPerBlock > input_desc[0].dims.d[1]) {
         cPerBlock = 8;
       }
+      auto d_dim = input_desc[0].dims.nbDims;
+      params_.n = input_desc[0].dims.d[0];
+      if (d_dim == 3) {
+        params_.c = input_desc[0].dims.d[1];
+        params_.d = 1;
+        params_.h = 1;
+        params_.w = input_desc[0].dims.d[2];
+      } else if (d_dim == 4) {
+        params_.c = input_desc[0].dims.d[1];
+        params_.d = 1;
+        params_.h = input_desc[0].dims.d[2];
+        params_.w = input_desc[0].dims.d[3];
+      } else {
+        // d_dim == 5
+        params_.c = input_desc[0].dims.d[1];
+        params_.d = input_desc[0].dims.d[2];
+        params_.h = input_desc[0].dims.d[3];
+        params_.w = input_desc[0].dims.d[4];
+      }
       params_.withSilu = with_silu_;
       params_.dst = static_cast<half *>(outputs[0]);
       params_.srcX = static_cast<half const *>(inputs[0]);
 
       params_.gamma = scale_gpu_;
       params_.beta = bias_gpu_;
       params_.redBuffer = static_cast<float *>(workspace);
-      params_.n = input_desc[0].dims.d[0];
-      params_.h = input_desc[0].dims.d[2];
-      params_.w = input_desc[0].dims.d[3];
-      params_.c = input_desc[0].dims.d[1];
+      // params_.n = input_desc[0].dims.d[0];
+      // params_.h = input_desc[0].dims.d[2];
+      // params_.w = input_desc[0].dims.d[3];
+      // params_.c = input_desc[0].dims.d[1];
       params_.groups = groups_;
-      params_.hw = params_.h * params_.w;
-      const int32_t blocksPerHW = findMaxDivisor(params_.hw, maxBlocksPerHW);
-      params_.hwPerBlock = divUp(params_.hw, blocksPerHW);
+      params_.dhw = params_.d * params_.h * params_.w;
+      const int32_t blocksPerDHW = findMaxDivisor(params_.dhw, maxBlocksPerDHW);
+      params_.dhwPerBlock = divUp(params_.dhw, blocksPerDHW);
       params_.cPerBlock = cPerBlock;
       params_.cPerGroup = params_.c / params_.groups;
-      params_.hwc = params_.hw * params_.c;
-      params_.invHWC = 1.F / static_cast<float>(params_.hw * params_.cPerGroup);
+      params_.dhwc = params_.dhw * params_.c;
+      params_.invDHWC =
+          1.F / static_cast<float>(params_.dhw * params_.cPerGroup);
       params_.groupsPerBlock = cPerBlock / params_.cPerGroup;
       CHECK_EQ(cPerBlock % params_.cPerGroup, 0);
       CHECK_EQ(params_.cPerGroup % 2, 0);
 
@@ -29,7 +29,7 @@ namespace inference {
 namespace tensorrt {
 namespace plugin {
 
-using phi::GroupNormNHWCParams;
+using phi::GroupNormNDHWCParams;
 class GroupNormPlugin : public PluginTensorRT {
  public:
   size_t getSerializationSize() const TRT_NOEXCEPT override {
@@ -289,7 +289,7 @@ class GroupNormPluginDynamic : public DynamicPluginTensorRT {
   float eps_;
   std::vector<int64_t> mean_shape_;
   std::vector<int64_t> variance_shape_;
-  GroupNormNHWCParams<half> params_;
+  GroupNormNDHWCParams<half> params_;
   bool with_silu_;
   bool with_fp16_;
   bool with_int8_;