[PIR] fix unique in pir (#67931)

wanghuancoder · web-flow · commit 87fbd08586d1 · 2024-09-03T15:51:09.000+08:00
* fix unique in pir

* refine

* refine
diff --git a/paddle/phi/kernels/gpu/unique_consecutive_functor.h b/paddle/phi/kernels/gpu/unique_consecutive_functor.h
@@ -174,6 +174,11 @@ static void ComputeUniqueConsecutiveDims(const Context& context,
                                          DenseTensor* inverse,
                                          DenseTensor* counts) {
   // 1. inverse indices: 'inverse'
+  DenseTensor tmp;
+  if (!inverse) {
+    inverse = &tmp;
+  }
+
   inverse->Resize(common::make_ddim({row}));
   auto inverse_data = context.template Alloc<IndexT>(inverse);
   DenseTensor inv_loc;
@@ -214,11 +219,15 @@ static void ComputeUniqueConsecutiveDims(const Context& context,
   sorted_indices->Resize(common::make_ddim({num_out}));
 
   // 3. counts: 'counts'
-  counts->Resize(common::make_ddim({num_out}));
-  auto count_data = context.template Alloc<IndexT>(counts);
-  thrust::fill(thrust::device, count_data, count_data + row, 0);
-  thrust::adjacent_difference(
-      thrust::device, range_data_ptr + 1, range_data_ptr + row + 1, count_data);
+  if (return_counts) {
+    counts->Resize(common::make_ddim({num_out}));
+    auto count_data = context.template Alloc<IndexT>(counts);
+    thrust::fill(thrust::device, count_data, count_data + row, 0);
+    thrust::adjacent_difference(thrust::device,
+                                range_data_ptr + 1,
+                                range_data_ptr + row + 1,
+                                count_data);
+  }
 }
 
 // Binary function 'equal_to'
diff --git a/paddle/phi/kernels/gpu/unique_kernel.cu b/paddle/phi/kernels/gpu/unique_kernel.cu
@@ -127,6 +127,10 @@ UniqueFlattendCUDATensor(const Context& context,
   DenseTensor in_hat;
   phi::Copy(context, in, context.GetPlace(), false, &in_hat);
   auto* in_data_hat = context.template Alloc<InT>(&in_hat);
+  DenseTensor tmp;
+  if (!indices) {
+    indices = &tmp;
+  }
 
   indices->Resize(common::make_ddim({num_input}));
   auto* indices_data = context.template Alloc<IndexT>(indices);
@@ -261,6 +265,11 @@ UniqueFlattendCUDATensor(const Context& context,
   auto equal = BinaryEqual<InT>(1, in_data);
   auto not_equal = BinaryNotEqual<InT>(1, in_data);
 
+  DenseTensor tmp;
+  if (!indices) {
+    indices = &tmp;
+  }
+
   indices->Resize(common::make_ddim({num_input}));
   auto* indices_data = context.template Alloc<IndexT>(indices);
 
@@ -400,13 +409,15 @@ static void ComputeUniqueDims(const Context& context,
   sorted_indices->Resize(common::make_ddim({num_out}));
 
   // 3. counts: 'counts'
-  counts->Resize(common::make_ddim({num_out}));
-  auto* count_data = context.template Alloc<IndexT>(counts);
-  thrust::fill(exec_policy, count_data, count_data + num_out, 0);
-  thrust::adjacent_difference(exec_policy,
-                              range_data_ptr + 1,
-                              range_data_ptr + num_out + 1,
-                              count_data);
+  if (return_counts) {
+    counts->Resize(common::make_ddim({num_out}));
+    auto* count_data = context.template Alloc<IndexT>(counts);
+    thrust::fill(exec_policy, count_data, count_data + num_out, 0);
+    thrust::adjacent_difference(exec_policy,
+                                range_data_ptr + 1,
+                                range_data_ptr + num_out + 1,
+                                count_data);
+  }
 }
 
 // Calculate unique when 'axis' is set
@@ -455,6 +466,11 @@ static void UniqueDimsCUDATensor(const Context& context,
   int64_t row = in_trans.dims()[0];
   const InT* in_trans_data = in_trans.data<InT>();
 
+  DenseTensor tmp;
+  if (!indices) {
+    indices = &tmp;
+  }
+
   indices->Resize(common::make_ddim({row}));
   auto* sorted_indices_data = context.template Alloc<IndexT>(indices);
 
