some kernel support int64 input (#69829)

wanghuancoder · web-flow · commit 813999530224 · 2024-12-02T14:53:55.000+08:00
* some kernel support int64 input
diff --git a/paddle/phi/kernels/cpu/linspace_kernel.cc b/paddle/phi/kernels/cpu/linspace_kernel.cc
@@ -27,7 +27,12 @@ void LinspaceKernel(const Context& ctx,
                     const DenseTensor& number,
                     DataType dtype,
                     DenseTensor* out) {
-  int32_t num = number.data<int32_t>()[0];
+  int64_t num = 0;
+  if (number.dtype() == phi::DataType::INT64) {
+    num = number.data<int64_t>()[0];
+  } else if (number.dtype() == phi::DataType::INT32) {
+    num = number.data<int32_t>()[0];
+  }
   auto start_t = phi::funcs::TransDataType(ctx, start, dtype);
   auto stop_t = phi::funcs::TransDataType(ctx, stop, dtype);
 
diff --git a/paddle/phi/kernels/cpu/multiclass_nms3_kernel.cc b/paddle/phi/kernels/cpu/multiclass_nms3_kernel.cc
@@ -236,10 +236,20 @@ T PolyIoU(const T* box1,
 
 inline std::vector<size_t> GetNmsLodFromRoisNum(const DenseTensor* rois_num) {
   std::vector<size_t> rois_lod;
-  auto* rois_num_data = rois_num->data<int>();
-  rois_lod.push_back(static_cast<size_t>(0));
-  for (int i = 0; i < rois_num->numel(); ++i) {
-    rois_lod.push_back(rois_lod.back() + static_cast<size_t>(rois_num_data[i]));
+  if (rois_num->dtype() == phi::DataType::INT64) {
+    auto* rois_num_data = rois_num->data<int64_t>();
+    rois_lod.push_back(static_cast<size_t>(0));
+    for (int64_t i = 0; i < rois_num->numel(); ++i) {
+      rois_lod.push_back(rois_lod.back() +
+                         static_cast<size_t>(rois_num_data[i]));
+    }
+  } else if (rois_num->dtype() == phi::DataType::INT32) {
+    auto* rois_num_data = rois_num->data<int>();
+    rois_lod.push_back(static_cast<size_t>(0));
+    for (int i = 0; i < rois_num->numel(); ++i) {
+      rois_lod.push_back(rois_lod.back() +
+                         static_cast<size_t>(rois_num_data[i]));
+    }
   }
   return rois_lod;
 }
diff --git a/paddle/phi/kernels/cpu/roi_align_kernel.cc b/paddle/phi/kernels/cpu/roi_align_kernel.cc
@@ -217,13 +217,24 @@ void RoiAlignKernel(const Context& dev_ctx,
             "and the batch size of images is %d",
             boxes_batch_size,
             batch_size));
-    auto* boxes_num_data = boxes_num->data<int>();
-    int start = 0;
-    for (int n = 0; n < boxes_batch_size; ++n) {
-      for (int i = start; i < start + boxes_num_data[n]; ++i) {
-        roi_batch_id_data[i] = n;
+    if (boxes_num->dtype() == phi::DataType::INT64) {
+      auto* boxes_num_data = boxes_num->data<int64_t>();
+      int64_t start = 0;
+      for (int64_t n = 0; n < boxes_batch_size; ++n) {
+        for (int64_t i = start; i < start + boxes_num_data[n]; ++i) {
+          roi_batch_id_data[i] = n;
+        }
+        start += boxes_num_data[n];
+      }
+    } else if (boxes_num->dtype() == phi::DataType::INT32) {
+      auto* boxes_num_data = boxes_num->data<int>();
+      int start = 0;
+      for (int n = 0; n < boxes_batch_size; ++n) {
+        for (int i = start; i < start + boxes_num_data[n]; ++i) {
+          roi_batch_id_data[i] = n;
+        }
+        start += boxes_num_data[n];
       }
-      start += boxes_num_data[n];
     }
   } else {
     auto lod = boxes.lod();
diff --git a/paddle/phi/kernels/funcs/distribute_fpn_proposals_functor.h b/paddle/phi/kernels/funcs/distribute_fpn_proposals_functor.h
@@ -33,16 +33,32 @@ template <typename Context>
 inline std::vector<size_t> GetLodFromRoisNum(const Context& dev_ctx,
                                              const DenseTensor* rois_num) {
   std::vector<size_t> rois_lod;
-  auto* rois_num_data = rois_num->data<int>();
-  DenseTensor cpu_tensor;
-  if (rois_num->place().GetType() == phi::AllocationType::GPU ||
-      rois_num->place().GetType() == phi::AllocationType::XPU) {
-    Copy<Context>(dev_ctx, *rois_num, phi::CPUPlace(), true, &cpu_tensor);
-    rois_num_data = cpu_tensor.data<int>();
-  }
-  rois_lod.push_back(static_cast<size_t>(0));
-  for (int i = 0; i < rois_num->numel(); ++i) {
-    rois_lod.push_back(rois_lod.back() + static_cast<size_t>(rois_num_data[i]));
+  if (rois_num->dtype() == phi::DataType::INT64) {
+    auto* rois_num_data = rois_num->data<int64_t>();
+    DenseTensor cpu_tensor;
+    if (rois_num->place().GetType() == phi::AllocationType::GPU ||
+        rois_num->place().GetType() == phi::AllocationType::XPU) {
+      Copy<Context>(dev_ctx, *rois_num, phi::CPUPlace(), true, &cpu_tensor);
+      rois_num_data = cpu_tensor.data<int64_t>();
+    }
+    rois_lod.push_back(static_cast<size_t>(0));
+    for (int64_t i = 0; i < rois_num->numel(); ++i) {
+      rois_lod.push_back(rois_lod.back() +
+                         static_cast<size_t>(rois_num_data[i]));
+    }
+  } else if (rois_num->dtype() == phi::DataType::INT32) {
+    auto* rois_num_data = rois_num->data<int>();
+    DenseTensor cpu_tensor;
+    if (rois_num->place().GetType() == phi::AllocationType::GPU ||
+        rois_num->place().GetType() == phi::AllocationType::XPU) {
+      Copy<Context>(dev_ctx, *rois_num, phi::CPUPlace(), true, &cpu_tensor);
+      rois_num_data = cpu_tensor.data<int>();
+    }
+    rois_lod.push_back(static_cast<size_t>(0));
+    for (int i = 0; i < rois_num->numel(); ++i) {
+      rois_lod.push_back(rois_lod.back() +
+                         static_cast<size_t>(rois_num_data[i]));
+    }
   }
   return rois_lod;
 }
diff --git a/paddle/phi/kernels/funcs/interpolate_function.h b/paddle/phi/kernels/funcs/interpolate_function.h
@@ -94,28 +94,54 @@ inline std::vector<int> get_new_shape(
                           "The shape of dimension tensor should be [1] or [],"
                           "but received d%.",
                           tensor->dims()));
+    if (tensor->dtype() == phi::DataType::INT64) {
 #ifdef PADDLE_WITH_CUSTOM_DEVICE
-    if (tensor->place().GetType() == phi::AllocationType::CUSTOM) {
-      DenseTensor temp;
-      phi::Copy(*dev_ctx, *tensor, phi::CPUPlace(), true, &temp);
-      vec_new_shape.push_back(static_cast<int32_t>(*temp.data<int32_t>()));
-      continue;
-    }
+      if (tensor->place().GetType() == phi::AllocationType::CUSTOM) {
+        DenseTensor temp;
+        phi::Copy(*dev_ctx, *tensor, phi::CPUPlace(), true, &temp);
+        vec_new_shape.push_back(static_cast<int64_t>(*temp.data<int64_t>()));
+        continue;
+      }
 #endif
 #ifdef PADDLE_WITH_XPU
-    if (tensor->place().GetType() == phi::AllocationType::XPU) {
-      DenseTensor temp;
-      phi::Copy(*dev_ctx, *tensor, phi::CPUPlace(), true, &temp);
-      vec_new_shape.push_back(static_cast<int32_t>(*temp.data<int32_t>()));
-      continue;
-    }
+      if (tensor->place().GetType() == phi::AllocationType::XPU) {
+        DenseTensor temp;
+        phi::Copy(*dev_ctx, *tensor, phi::CPUPlace(), true, &temp);
+        vec_new_shape.push_back(static_cast<int64_t>(*temp.data<int64_t>()));
+        continue;
+      }
+#endif
+      if (tensor->place().GetType() == phi::AllocationType::GPU) {
+        DenseTensor temp;
+        phi::Copy(*dev_ctx, *tensor, phi::CPUPlace(), true, &temp);
+        vec_new_shape.push_back(static_cast<int64_t>(*temp.data<int64_t>()));
+      } else {
+        vec_new_shape.push_back(static_cast<int64_t>(*tensor->data<int64_t>()));
+      }
+    } else if (tensor->dtype() == phi::DataType::INT32) {
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+      if (tensor->place().GetType() == phi::AllocationType::CUSTOM) {
+        DenseTensor temp;
+        phi::Copy(*dev_ctx, *tensor, phi::CPUPlace(), true, &temp);
+        vec_new_shape.push_back(static_cast<int32_t>(*temp.data<int32_t>()));
+        continue;
+      }
+#endif
+#ifdef PADDLE_WITH_XPU
+      if (tensor->place().GetType() == phi::AllocationType::XPU) {
+        DenseTensor temp;
+        phi::Copy(*dev_ctx, *tensor, phi::CPUPlace(), true, &temp);
+        vec_new_shape.push_back(static_cast<int32_t>(*temp.data<int32_t>()));
+        continue;
+      }
 #endif
-    if (tensor->place().GetType() == phi::AllocationType::GPU) {
-      DenseTensor temp;
-      phi::Copy(*dev_ctx, *tensor, phi::CPUPlace(), true, &temp);
-      vec_new_shape.push_back(static_cast<int32_t>(*temp.data<int32_t>()));
-    } else {
-      vec_new_shape.push_back(static_cast<int32_t>(*tensor->data<int32_t>()));
+      if (tensor->place().GetType() == phi::AllocationType::GPU) {
+        DenseTensor temp;
+        phi::Copy(*dev_ctx, *tensor, phi::CPUPlace(), true, &temp);
+        vec_new_shape.push_back(static_cast<int32_t>(*temp.data<int32_t>()));
+      } else {
+        vec_new_shape.push_back(static_cast<int32_t>(*tensor->data<int32_t>()));
+      }
     }
   }
 
diff --git a/paddle/phi/kernels/gpu/roi_align_kernel.cu b/paddle/phi/kernels/gpu/roi_align_kernel.cu
@@ -181,19 +181,36 @@ void RoiAlignKernel(const Context& dev_ctx,
             boxes_batch_size,
             batch_size));
 
-    std::vector<int> boxes_num_list(boxes_batch_size);
-    memory_utils::Copy(cplace,
-                       boxes_num_list.data(),
-                       gplace,
-                       boxes_num->data<int>(),
-                       sizeof(int) * boxes_batch_size,
-                       0);
-    int start = 0;
-    for (int n = 0; n < boxes_batch_size; ++n) {
-      for (int i = start; i < start + boxes_num_list[n]; ++i) {
-        roi_batch_id_data[i] = n;
+    if (boxes_num->dtype() == phi::DataType::INT64) {
+      std::vector<int64_t> boxes_num_list(boxes_batch_size);
+      memory_utils::Copy(cplace,
+                         boxes_num_list.data(),
+                         gplace,
+                         boxes_num->data<int64_t>(),
+                         sizeof(int64_t) * boxes_batch_size,
+                         0);
+      int64_t start = 0;
+      for (int64_t n = 0; n < boxes_batch_size; ++n) {
+        for (int64_t i = start; i < start + boxes_num_list[n]; ++i) {
+          roi_batch_id_data[i] = n;
+        }
+        start += boxes_num_list[n];
+      }
+    } else if (boxes_num->dtype() == phi::DataType::INT32) {
+      std::vector<int> boxes_num_list(boxes_batch_size);
+      memory_utils::Copy(cplace,
+                         boxes_num_list.data(),
+                         gplace,
+                         boxes_num->data<int>(),
+                         sizeof(int) * boxes_batch_size,
+                         0);
+      int start = 0;
+      for (int n = 0; n < boxes_batch_size; ++n) {
+        for (int i = start; i < start + boxes_num_list[n]; ++i) {
+          roi_batch_id_data[i] = n;
+        }
+        start += boxes_num_list[n];
       }
-      start += boxes_num_list[n];
     }
   } else {
     auto lod = boxes.lod();
diff --git a/paddle/phi/kernels/xpu/linspace_kernel.cc b/paddle/phi/kernels/xpu/linspace_kernel.cc
@@ -58,7 +58,7 @@ void LinspaceKernel(const Context& ctx,
   using XPUType = typename XPUTypeTrait<T>::Type;
   T start_value = GetValueOfExpectedType<T, Context>(ctx, start);
   T stop_value = GetValueOfExpectedType<T, Context>(ctx, stop);
-  int32_t num = GetValueOfExpectedType<int32_t, Context>(ctx, number);
+  int64_t num = GetValueOfExpectedType<int64_t, Context>(ctx, number);
 
   PADDLE_ENFORCE_GT(
       num,
diff --git a/paddle/phi/kernels/xpu/multiclass_nms3_kernel.cc b/paddle/phi/kernels/xpu/multiclass_nms3_kernel.cc
@@ -63,16 +63,30 @@ void MultiClassNMSKernel(const Context& ctx,
     if (has_rois_num) {
       phi::DenseTensor rois_num_host;
       rois_num_host.Resize(rois_num.get_ptr()->dims());
-      ctx.template HostAlloc<int>(&rois_num_host);
-      phi::Copy(ctx,
-                *rois_num.get_ptr(),
-                rois_num_host.place(),
-                false,
-                &rois_num_host);
-      n = rois_num.get_ptr()->numel();
-      for (int i = 0; i < n; i++) {
-        rois_num_vec.push_back(rois_num_host.data<int>()[i]);
-        boxes_count += rois_num_host.data<int>()[i];
+      if (rois_num.get_ptr()->dtype() == phi::DataType::INT64) {
+        ctx.template HostAlloc<int64_t>(&rois_num_host);
+        phi::Copy(ctx,
+                  *rois_num.get_ptr(),
+                  rois_num_host.place(),
+                  false,
+                  &rois_num_host);
+        n = rois_num.get_ptr()->numel();
+        for (int64_t i = 0; i < n; i++) {
+          rois_num_vec.push_back(rois_num_host.data<int64_t>()[i]);
+          boxes_count += rois_num_host.data<int64_t>()[i];
+        }
+      } else if (rois_num.get_ptr()->dtype() == phi::DataType::INT32) {
+        ctx.template HostAlloc<int>(&rois_num_host);
+        phi::Copy(ctx,
+                  *rois_num.get_ptr(),
+                  rois_num_host.place(),
+                  false,
+                  &rois_num_host);
+        n = rois_num.get_ptr()->numel();
+        for (int i = 0; i < n; i++) {
+          rois_num_vec.push_back(rois_num_host.data<int>()[i]);
+          boxes_count += rois_num_host.data<int>()[i];
+        }
       }
     } else {
       auto lod = bboxes.lod().back();
diff --git a/paddle/phi/kernels/xpu/roi_align_kernel.cc b/paddle/phi/kernels/xpu/roi_align_kernel.cc
@@ -64,16 +64,30 @@ void RoiAlignKernel(const Context& dev_ctx,
             rois_batch_size,
             batch_size));
 
-    std::vector<int> rois_num_list(rois_batch_size);
-    memory_utils::Copy(cplace,
-                       rois_num_list.data(),
-                       xplace,
-                       boxes_num->data<int>(),
-                       sizeof(int) * rois_batch_size);
-    cpu_lod = new int[rois_batch_size + 1];
-    cpu_lod[0] = 0;
-    for (int i = 0; i < rois_batch_size; i++) {
-      cpu_lod[i + 1] = cpu_lod[i] + rois_num_list[i];
+    if (boxes_num->dtype() == phi::DataType::INT64) {
+      std::vector<int64_t> rois_num_list(rois_batch_size);
+      memory_utils::Copy(cplace,
+                         rois_num_list.data(),
+                         xplace,
+                         boxes_num->data<int64_t>(),
+                         sizeof(int64_t) * rois_batch_size);
+      cpu_lod = new int[rois_batch_size + 1];
+      cpu_lod[0] = 0;
+      for (int64_t i = 0; i < rois_batch_size; i++) {
+        cpu_lod[i + 1] = cpu_lod[i] + rois_num_list[i];
+      }
+    } else if (boxes_num->dtype() == phi::DataType::INT32) {
+      std::vector<int> rois_num_list(rois_batch_size);
+      memory_utils::Copy(cplace,
+                         rois_num_list.data(),
+                         xplace,
+                         boxes_num->data<int>(),
+                         sizeof(int) * rois_batch_size);
+      cpu_lod = new int[rois_batch_size + 1];
+      cpu_lod[0] = 0;
+      for (int i = 0; i < rois_batch_size; i++) {
+        cpu_lod[i + 1] = cpu_lod[i] + rois_num_list[i];
+      }
     }
   } else {
     auto lod = boxes.lod();
diff --git a/python/paddle/tensor/creation.py b/python/paddle/tensor/creation.py
@@ -396,7 +396,7 @@ def linspace(
         else:
             check_type(stop, 'stop', (int, float), 'linspace')
         if isinstance(num, paddle.pir.Value):
-            check_dtype(num.dtype, 'num', ['int32'], 'linspace')
+            check_dtype(num.dtype, 'num', ['int32', 'int64'], 'linspace')
         check_dtype(
             dtype,
             'dtype',
