Adapt new dipu device (#2905)

open-mmlab · Aug 22, 2023 · ee93530 · ee93530
1 parent c523359
commit ee93530
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 7 deletions.
diff --git a/mmcv/ops/csrc/pytorch/bbox_overlaps.cpp b/mmcv/ops/csrc/pytorch/bbox_overlaps.cpp
@@ -6,6 +6,7 @@
 #include <diopi/functions.h>
 #include <diopi/functions_mmcv.h>
 
+#include "csrc_dipu/base/basedef.h"
 #include "csrc_dipu/diopirt/diopirt_impl.h"
 
 using dipu::diopi_helper::toDiopiScalar;
@@ -33,7 +34,7 @@ void bbox_overlaps_diopi(const Tensor bboxes1, const Tensor bboxes2,
   diopiContextHandle_t ch = &ctx;
   auto bboxes2_p = toDiopiTensorHandle(bboxes2);
   auto ious_p = toDiopiTensorHandle(ious);
-  bool is_mock_cuda = bboxes1.device().type() == c10::DeviceType::PrivateUse1;
+  bool is_mock_cuda = bboxes1.device().type() == dipu::DIPU_DEVICE_TYPE;
   if (is_mock_cuda &&
       reinterpret_cast<void *>(diopiBboxOverlapsMmcv) != nullptr) {
     auto ret = diopiBboxOverlapsMmcv(ch, ious_p, bboxes1_p, bboxes2_p, mode,

diff --git a/mmcv/ops/csrc/pytorch/nms.cpp b/mmcv/ops/csrc/pytorch/nms.cpp
@@ -6,6 +6,7 @@
 #include <diopi/functions.h>
 #include <diopi/functions_mmcv.h>
 
+#include "csrc_dipu/base/basedef.h"
 #include "csrc_dipu/diopirt/diopirt_impl.h"
 
 using dipu::diopi_helper::toDiopiScalar;
@@ -42,7 +43,7 @@ Tensor nms_diopi(Tensor boxes, Tensor scores, float iou_threshold, int offset) {
   auto outp = toDiopiTensorHandle(out);
   diopiTensorHandle_t* outhandle = &outp;
   auto scores_p = toDiopiTensorHandle(scores);
-  bool is_mock_cuda = boxes.device().type() == c10::DeviceType::PrivateUse1;
+  bool is_mock_cuda = boxes.device().type() == dipu::DIPU_DEVICE_TYPE;
   if (is_mock_cuda && reinterpret_cast<void*>(diopiNmsMmcv) != nullptr) {
     auto ret =
         diopiNmsMmcv(ch, outhandle, boxes_p, scores_p, iou_threshold, offset);

diff --git a/mmcv/ops/csrc/pytorch/roi_align.cpp b/mmcv/ops/csrc/pytorch/roi_align.cpp
@@ -6,6 +6,7 @@
 #include <diopi/functions.h>
 #include <diopi/functions_mmcv.h>
 
+#include "csrc_dipu/base/basedef.h"
 #include "csrc_dipu/diopirt/diopirt_impl.h"
 
 using dipu::diopi_helper::toDiopiScalar;
@@ -53,8 +54,8 @@ void roi_align_forward_diopi(Tensor input, Tensor rois, Tensor output,
   auto out_p = toDiopiTensorHandle(output);
   auto argmax_y_p = toDiopiTensorHandle(argmax_y);
   auto argmax_x_p = toDiopiTensorHandle(argmax_x);
-  bool is_mock_cuda = input.device().type() == c10::DeviceType::PrivateUse1;
-  if (is_mock_cuda && reinterpret_cast<void*>(diopiRoiAlignMmcv) != nullptr) {
+  bool is_mock_cuda = input.device().type() == dipu::DIPU_DEVICE_TYPE;
+  if (is_mock_cuda && reinterpret_cast<void *>(diopiRoiAlignMmcv) != nullptr) {
     auto ret = diopiRoiAlignMmcv(
         ch, out_p, argmax_y_p, argmax_x_p, input_p, rois_p, aligned_height,
         aligned_width, sampling_ratio, pool_mode, spatial_scale, aligned);
@@ -92,10 +93,9 @@ void roi_align_backward_diopi(Tensor grad_output, Tensor rois, Tensor argmax_y,
   auto grad_input_ = toDiopiTensorHandle(grad_input);
   diopiContext ctx(dipu::getCurrentDIPUStream().rawstream());
   diopiContextHandle_t ch = &ctx;
-  bool is_mock_cuda =
-      grad_output.device().type() == c10::DeviceType::PrivateUse1;
+  bool is_mock_cuda = grad_output.device().type() == dipu::DIPU_DEVICE_TYPE;
   if (is_mock_cuda &&
-      reinterpret_cast<void*>(diopiRoiAlignBackwardMmcv) != nullptr) {
+      reinterpret_cast<void *>(diopiRoiAlignBackwardMmcv) != nullptr) {
     auto ret = diopiRoiAlignBackwardMmcv(ch, grad_input_, grad_output_, rois_,
                                          argmax_y_, argmax_x_, aligned_height,
                                          aligned_width, sampling_ratio,