make multiclass_nms3 kernel support hip (#65599)

Deleter-D · web-flow · commit 04c38b05c5b2 · 2024-07-02T10:27:02.000+08:00
diff --git a/paddle/phi/kernels/gpu/multiclass_nms3_kernel.cu b/paddle/phi/kernels/gpu/multiclass_nms3_kernel.cu
@@ -12,11 +12,14 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License. */
 
-#ifndef PADDLE_WITH_HIP
-
 #include "paddle/phi/kernels/multiclass_nms3_kernel.h"
 
+#ifdef PADDLE_WITH_HIP
+#include <hipcub/hipcub.hpp>
+namespace cub = hipcub;
+#else
 #include <cub/cub.cuh>
+#endif
 #include "cuda.h"  // NOLINT
 
 #include "paddle/phi/backends/context_pool.h"
@@ -234,7 +237,7 @@ __launch_bounds__(nthds_per_cta) __global__
 }
 
 template <typename T_SCORE>
-void SortScoresPerClassGPU(cudaStream_t stream,
+void SortScoresPerClassGPU(gpuStream_t stream,
                            const int num,
                            const int num_classes,
                            const int num_preds_per_class,
@@ -298,7 +301,11 @@ void SortScoresPerClassGPU(cudaStream_t stream,
       begin_bit,
       end_bit,
       stream);
+#ifdef PADDLE_WITH_HIP
+  PADDLE_ENFORCE_GPU_SUCCESS(hipGetLastError());
+#else
   PADDLE_ENFORCE_GPU_SUCCESS(cudaGetLastError());
+#endif
 }
 
 /* ===========
@@ -539,7 +546,7 @@ __global__ void AllClassNMSKernel(
 }
 
 template <typename T_SCORE, typename T_BBOX>
-void AllClassNMSGPU(cudaStream_t stream,
+void AllClassNMSGPU(gpuStream_t stream,
                     const int num,
                     const int num_classes,
                     const int num_preds_per_class,
@@ -603,7 +610,11 @@ void AllClassNMSGPU(cudaStream_t stream,
       score_shift,
       caffe_semantics);
 
+#ifdef PADDLE_WITH_HIP
+  PADDLE_ENFORCE_GPU_SUCCESS(hipGetLastError());
+#else
   PADDLE_ENFORCE_GPU_SUCCESS(cudaGetLastError());
+#endif
 }
 
 /* ==================
@@ -618,7 +629,7 @@ __launch_bounds__(nthds_per_cta) __global__
   if (idx <= num_segments) d_offsets[idx] = idx * offset;
 }
 
-void SetUniformOffsets(cudaStream_t stream,
+void SetUniformOffsets(gpuStream_t stream,
                        const int num_segments,
                        const int offset,
                        int* d_offsets) {
@@ -706,7 +717,7 @@ __launch_bounds__(nthds_per_cta) __global__
 }
 
 template <typename T_BBOX, typename T_SCORE>
-void GatherNMSOutputsGPU(cudaStream_t stream,
+void GatherNMSOutputsGPU(gpuStream_t stream,
                          const bool share_location,
                          const int num_images,
                          const int num_preds_per_class,
@@ -724,8 +735,13 @@ void GatherNMSOutputsGPU(cudaStream_t stream,
                          void* nmsed_valid_mask,
                          bool clip_boxes,
                          const float score_shift) {
+#ifdef PADDLE_WITH_HIP
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      hipMemsetAsync(num_detections, 0, num_images * sizeof(int), stream));
+#else
   PADDLE_ENFORCE_GPU_SUCCESS(
       cudaMemsetAsync(num_detections, 0, num_images * sizeof(int), stream));
+#endif
   const int BS = 32;
   const int GS = 32;
   GatherNMSOutputsKernel<T_BBOX, T_SCORE, BS>
@@ -746,12 +762,15 @@ void GatherNMSOutputsGPU(cudaStream_t stream,
                               reinterpret_cast<int*>(nmsed_valid_mask),
                               clip_boxes,
                               T_SCORE(score_shift));
-
+#ifdef PADDLE_WITH_HIP
+  PADDLE_ENFORCE_GPU_SUCCESS(hipGetLastError());
+#else
   PADDLE_ENFORCE_GPU_SUCCESS(cudaGetLastError());
+#endif
 }
 
 template <typename T_SCORE>
-void SortScoresPerImageGPU(cudaStream_t stream,
+void SortScoresPerImageGPU(gpuStream_t stream,
                            const int num_images,
                            const int num_items_per_image,
                            void* unsorted_scores,
@@ -792,11 +811,15 @@ void SortScoresPerImageGPU(cudaStream_t stream,
       begin_bit,
       end_bit,
       stream);
+#ifdef PADDLE_WITH_HIP
+  PADDLE_ENFORCE_GPU_SUCCESS(hipGetLastError());
+#else
   PADDLE_ENFORCE_GPU_SUCCESS(cudaGetLastError());
+#endif
 }
 
 template <typename T>
-void InferNMS(cudaStream_t stream,
+void InferNMS(gpuStream_t stream,
               const int N,
               const int per_batch_boxes_size,
               const int per_batch_scores_size,
@@ -831,20 +854,33 @@ void InferNMS(cudaStream_t stream,
   size_t bbox_data_size =
       CalcDetectionForwardBBoxDataSize<T>(N, per_batch_boxes_size);
   void* bbox_data_raw = workspace;
+#ifdef PADDLE_WITH_HIP
+  PADDLE_ENFORCE_GPU_SUCCESS(hipMemcpyAsync(bbox_data_raw,
+                                            loc_data,
+                                            bbox_data_size,
+                                            hipMemcpyDeviceToDevice,
+                                            stream));
+#else
   PADDLE_ENFORCE_GPU_SUCCESS(cudaMemcpyAsync(bbox_data_raw,
                                              loc_data,
                                              bbox_data_size,
                                              cudaMemcpyDeviceToDevice,
                                              stream));
+#endif
   void* bbox_data = bbox_data_raw;
 
   const int num_scores = N * per_batch_scores_size;
   size_t total_scores_size =
       CalcDetectionForwardPreNMSSize<T>(N, per_batch_scores_size);
   void* scores =
       GetNextWorkspacePtr(reinterpret_cast<int8_t*>(bbox_data), bbox_data_size);
+#ifdef PADDLE_WITH_HIP
+  PADDLE_ENFORCE_GPU_SUCCESS(hipMemcpyAsync(
+      scores, conf_data, total_scores_size, hipMemcpyDeviceToDevice, stream));
+#else
   PADDLE_ENFORCE_GPU_SUCCESS(cudaMemcpyAsync(
       scores, conf_data, total_scores_size, cudaMemcpyDeviceToDevice, stream));
+#endif
 
   size_t indices_size =
       CalcDetectionForwardPreNMSSize<int>(N, per_batch_scores_size);
@@ -1136,13 +1172,8 @@ void MultiClassNMSGPUKernel(const Context& ctx,
 
 }  // namespace phi
 
-PD_REGISTER_KERNEL(multiclass_nms3,  // cuda_only
-                   GPU,
-                   ALL_LAYOUT,
-                   phi::MultiClassNMSGPUKernel,
-                   float) {
+PD_REGISTER_KERNEL(
+    multiclass_nms3, GPU, ALL_LAYOUT, phi::MultiClassNMSGPUKernel, float) {
   kernel->OutputAt(1).SetDataType(phi::DataType::INT32);
   kernel->OutputAt(2).SetDataType(phi::DataType::INT32);
 }
-
-#endif
