[XPU] Add lod support to roi_align op. (#9274)

Author: wbn03

lite/kernels/xpu/roi_align_compute.cc

@@ -30,53 +30,67 @@ void RoiAlignCompute::Run() {
 
   auto* in = param.X;
   auto* rois = param.ROIs;
+  auto* rois_num = param.RoisNum;
   auto* out = param.Out;
   float spatial_scale = param.spatial_scale;
   int pooled_height = param.pooled_height;
   int pooled_width = param.pooled_width;
   int sampling_ratio = param.sampling_ratio;
+  bool align = param.align;
 
   auto in_dims = in->dims();
   int batch_size = in_dims[0];
   int channels = in_dims[1];
   int height = in_dims[2];
   int width = in_dims[3];
   auto rois_dims = rois->dims();
-  int rois_num = rois_dims[0];
+  int rois_total = rois_dims[0];
   auto out_dims = out->dims();
-  if (rois_num == 0) {
+  if (rois_total == 0) {
     return;
   }
-  auto rois_lod = rois->lod().back();
+  const int* xpu_lod = nullptr;
   std::vector<int> cpu_lod_data;
   cpu_lod_data.resize(batch_size + 1);
-  for (int i = 0; i < rois_lod.size(); i++) {
-    cpu_lod_data[i] = rois_lod[i];
+  if (rois_num == nullptr) {
+    auto rois_lod = rois->lod().back();
+    int n = rois_lod.size();
+    for (int i = 0; i < n; i++) {
+      cpu_lod_data[i] = rois_lod[i];
+    }
+  } else {
+    auto rois_lod = rois_num->data<int>();
+    int n = rois_num->numel();
+    for (int i = 1; i <= n; i++) {
+      cpu_lod_data[i] = cpu_lod_data[i - 1] + rois_lod[i - 1];
+    }
   }
-
+  CHECK_EQ(cpu_lod_data[batch_size], rois_total);
   XPUScratchPadGuard xpu_lod_grad_ =
       TargetWrapperXPU::MallocScratchPad(cpu_lod_data.size() * sizeof(int));
-  int* xpu_lod = reinterpret_cast<int*>(xpu_lod_grad_->addr_);
-  XPU_CALL(xpu_memcpy(xpu_lod,
+  XPU_CALL(xpu_memcpy(xpu_lod_grad_->addr_,
                       cpu_lod_data.data(),
                       sizeof(int) * cpu_lod_data.size(),
                       XPUMemcpyKind::XPU_HOST_TO_DEVICE));
+  xpu_lod = reinterpret_cast<int*>(xpu_lod_grad_->addr_);
 
   int r = xdnn::roi_align<float, int>(ctx.GetRawContext(),
                                       in->data<float>(),
                                       out->mutable_data<float>(TARGET(kXPU)),
                                       rois->data<float>(),
-                                      static_cast<const int*>(xpu_lod),
+                                      xpu_lod,
                                       batch_size,
                                       channels,
                                       height,
                                       width,
-                                      rois_num,
+                                      rois_total,
                                       pooled_height,
                                       pooled_width,
                                       spatial_scale,
                                       sampling_ratio,
-                                      true);
+                                      true,
+                                      align,
+                                      false);
   CHECK_EQ(r, 0);
 }
 
@@ -93,5 +107,7 @@ REGISTER_LITE_KERNEL(roi_align,
                      def)
     .BindInput("X", {LiteType::GetTensorTy(TARGET(kXPU))})
     .BindInput("ROIs", {LiteType::GetTensorTy(TARGET(kXPU))})
+    .BindInput("RoisNum",
+               {LiteType::GetTensorTy(TARGET(kHost), PRECISION(kInt32))})
     .BindOutput("Out", {LiteType::GetTensorTy(TARGET(kXPU))})
     .Finalize();

lite/tests/kernels/roi_align_compute_test.cc

@@ -374,6 +374,8 @@ TEST(RoiAlign, precision) {
 #else
   return;
 #endif
+#elif defined(LITE_WITH_XPU)
+  place = TARGET(kXPU);
 #elif defined(LITE_WITH_X86) || defined(LITE_WITH_ARM)
   place = TARGET(kHost);
 #else