fix the unqiue op that generate the wrong the inreverse result (Paddl…

…ePaddle#62104)
hanhaowen-mt · Feb 27, 2024 · b89066a · b89066a
1 parent 867ab0d
commit b89066a
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/paddle/phi/kernels/gpu/unique_kernel.cu b/paddle/phi/kernels/gpu/unique_kernel.cu
@@ -168,8 +168,11 @@ UniqueFlattendCUDATensor(const Context& context,
 #ifdef PADDLE_WITH_HIP
     hipMemset(inv_loc_data_ptr, 0, sizeof(IndexT));
 #else
-    cudaMemsetAsync(inv_loc_data_ptr, 0, sizeof(IndexT), context.stream());
+    thrust::device_ptr<IndexT> inv_loc_data_dev(inv_loc_data_ptr);
+    inv_loc_data_dev[0] = 0;  // without device_ptr, segmentation fault
 #endif
+
+#ifdef PADDLE_WITH_HIP
     size_t temp_storage_bytes = 0;
     cub::DeviceScan::InclusiveSum(NULL,
                                   temp_storage_bytes,
@@ -185,6 +188,12 @@ UniqueFlattendCUDATensor(const Context& context,
                                   inv_loc_data_ptr,
                                   num_input,
                                   context.stream());
+#else
+    thrust::inclusive_scan(exec_policy,
+                           inv_loc_data_ptr,
+                           inv_loc_data_ptr + num_input,
+                           inv_loc_data_ptr);
+#endif
     thrust::scatter(exec_policy,
                     inv_loc_data_ptr,
                     inv_loc_data_ptr + num_input,