[Refactor] Replace tin_shift op of MLU backend with mlu-ops (#2911)

open-mmlab · Aug 28, 2023 · 2491dbb · 2491dbb
1 parent a0a1705
commit 2491dbb
Show file tree

Hide file tree

Showing 5 changed files with 76 additions and 444 deletions.
diff --git a/mmcv/ops/csrc/common/mlu/tin_shift_mlu_kernel.mlu b/mmcv/ops/csrc/common/mlu/tin_shift_mlu_kernel.mlu
diff --git a/mmcv/ops/csrc/pytorch/mlu/focal_loss_sigmoid_mlu.cpp b/mmcv/ops/csrc/pytorch/mlu/focal_loss_sigmoid_mlu.cpp
@@ -14,9 +14,9 @@
 
 #include "mlu_common_helper.h"
 
-void sigmoid_focal_loss_forward_mlu(Tensor input, Tensor target,
-                                    Tensor weight, Tensor output,
-                                    const float gamma, const float alpha) {
+void sigmoid_focal_loss_forward_mlu(Tensor input, Tensor target, Tensor weight,
+                                    Tensor output, const float gamma,
+                                    const float alpha) {
   // params check
   TORCH_CHECK(gamma >= 0, "gamma should be greater than or equal to 0. ",
               "But now gamma is ", gamma, ".");
@@ -82,15 +82,15 @@ void sigmoid_focal_loss_forward_mlu(Tensor input, Tensor target,
   auto handle = mluOpGetCurrentHandle();
 
   // launch kernel
-  TORCH_MLUOP_CHECK(mluOpFocalLossSigmoidForward(handle, prefer, reduction, input_desc.desc(),
-                               input_ptr, target_desc.desc(), target_ptr,
-                               weight_desc.desc(), weight_ptr, alpha, gamma,
-                               output_desc.desc(), output_ptr));
+  TORCH_MLUOP_CHECK(mluOpFocalLossSigmoidForward(
+      handle, prefer, reduction, input_desc.desc(), input_ptr,
+      target_desc.desc(), target_ptr, weight_desc.desc(), weight_ptr, alpha,
+      gamma, output_desc.desc(), output_ptr));
 }
 
-void sigmoid_focal_loss_backward_mlu(Tensor input, Tensor target,
-                                     Tensor weight, Tensor output,
-                                     const float gamma, const float alpha) {
+void sigmoid_focal_loss_backward_mlu(Tensor input, Tensor target, Tensor weight,
+                                     Tensor output, const float gamma,
+                                     const float alpha) {
   // params check
   TORCH_CHECK(gamma >= 0, "gamma should be greater than or equal to 0. ",
               "But now gamma is ", gamma, ".");
@@ -158,10 +158,10 @@ void sigmoid_focal_loss_backward_mlu(Tensor input, Tensor target,
   auto handle = mluOpGetCurrentHandle();
 
   // launch kernel
-  TORCH_MLUOP_CHECK(mluOpFocalLossSigmoidBackward(handle, prefer, reduction, input_desc.desc(),
-                                input_ptr, target_desc.desc(), target_ptr,
-                                weight_desc.desc(), weight_ptr, alpha, gamma,
-                                output_desc.desc(), output_ptr));
+  TORCH_MLUOP_CHECK(mluOpFocalLossSigmoidBackward(
+      handle, prefer, reduction, input_desc.desc(), input_ptr,
+      target_desc.desc(), target_ptr, weight_desc.desc(), weight_ptr, alpha,
+      gamma, output_desc.desc(), output_ptr));
 }
 
 void sigmoid_focal_loss_forward_impl(Tensor input, Tensor target, Tensor weight,

diff --git a/mmcv/ops/csrc/pytorch/mlu/mlu_common_helper.h b/mmcv/ops/csrc/pytorch/mlu/mlu_common_helper.h
@@ -18,8 +18,8 @@
 #include "pytorch_device_registry.hpp"
 
 #define MLUOP_MAJOR 0
-#define MLUOP_MINOR 7
-#define MLUOP_PATCHLEVEL 1
+#define MLUOP_MINOR 8
+#define MLUOP_PATCHLEVEL 0
 
 /*************************************************************************
  * This MACRO contains operations of simple tensor to mlu-tensor.