Add new parrots extension implementation for all ops (#794)

* delete all parrots file
add bbox_overlaps new parrots op impl

* support first new impl parrts op (bbox_overlaps)(success test)

* add box_iou_rotated op, test succeed

* add carafe and carafe_naive op, test succeed (one parrots bug need fix)

* add cc_attention op, test success

* add corner_pool op, test success

* add parrots op deform_conv, test success

* add deform_roi_pool op, test success (but has question)

* add focal loss op, test success (gradcheck)

* add masked_conv2d op, test success

* add modulated_deform_conv op, test success

* add nms and nms_rotated op, test success

* add psamask op, test success

* add roi_align op, test_success

* add roi_pool op, test success

* add sync_bn op, test success

* add tin_shift op, test success

* fix test_deform_roi_pool, add parrots test

* skip test_onnx because parrots does not support onnx

* fix c++ lint

* fix python lint

* fix python lint

mmcv/ops/csrc/parrots/cc_attention_parrots.cpp

+#include <parrots/compute/aten.hpp>
+#include <parrots/extension.hpp>
+#include <parrots/foundation/ssattrs.hpp>
+
+#include "cc_attention_pytorch.h"
+
+using namespace parrots;
+
+/*void ca_forward_cuda(const Tensor t, const Tensor f, Tensor weight);*/
+void ca_forward_cuda_parrots(CudaContext &ctx, const SSElement &attr,
+                             const OperatorBase::in_list_t &ins,
+                             OperatorBase::out_list_t &outs) {
+  const auto &t = buildATensor(ctx, ins[0]);
+  const auto &f = buildATensor(ctx, ins[1]);
+  auto weight = buildATensor(ctx, outs[0]);
+  ca_forward_cuda(t, f, weight);
+}
+
+/* void ca_backward_cuda(const Tensor dw, const Tensor t, const Tensor f,
+ *                       Tensor dt, Tensor df)
+ */
+void ca_backward_cuda_parrots(CudaContext &ctx, const SSElement &attr,
+                              const OperatorBase::in_list_t &ins,
+                              OperatorBase::out_list_t &outs) {
+  const auto &dw = buildATensor(ctx, ins[0]);
+  const auto &t = buildATensor(ctx, ins[1]);
+  const auto &f = buildATensor(ctx, ins[2]);
+  auto dt = buildATensor(ctx, outs[0]);
+  auto df = buildATensor(ctx, outs[1]);
+  ca_backward_cuda(dw, t, f, dt, df);
+}
+
+/* void ca_map_forward_cuda(const Tensor weight, const Tensor g, Tensor out); */
+void ca_map_forward_cuda_parrots(CudaContext &ctx, const SSElement &attr,
+                                 const OperatorBase::in_list_t &ins,
+                                 OperatorBase::out_list_t &outs) {
+  const auto &weight = buildATensor(ctx, ins[0]);
+  const auto &g = buildATensor(ctx, ins[1]);
+  auto out = buildATensor(ctx, outs[0]);
+  ca_map_forward_cuda(weight, g, out);
+}
+
+/* void ca_map_backward_cuda(const Tensor dout, const Tensor weight,
+ *                           const Tensor g, Tensor dw, Tensor dg);
+ */
+void ca_map_backward_cuda_parrots(CudaContext &ctx, const SSElement &attr,
+                                  const OperatorBase::in_list_t &ins,
+                                  OperatorBase::out_list_t &outs) {
+  const auto &dout = buildATensor(ctx, ins[0]);
+  const auto &weight = buildATensor(ctx, ins[1]);
+  const auto &g = buildATensor(ctx, ins[2]);
+  auto dw = buildATensor(ctx, outs[0]);
+  auto dg = buildATensor(ctx, outs[1]);
+  ca_map_backward_cuda(dout, weight, g, dw, dg);
+}
+
+PARROTS_EXTENSION_REGISTER(ca_forward)
+    .input(2)
+    .output(1)
+    .apply(ca_forward_cuda_parrots)
+    .done();
+
+PARROTS_EXTENSION_REGISTER(ca_backward)
+    .input(3)
+    .output(2)
+    .apply(ca_backward_cuda_parrots)
+    .done();
+
+PARROTS_EXTENSION_REGISTER(ca_map_forward)
+    .input(2)
+    .output(1)
+    .apply(ca_map_forward_cuda_parrots)
+    .done();
+
+PARROTS_EXTENSION_REGISTER(ca_map_backward)
+    .input(3)
+    .output(2)
+    .apply(ca_map_backward_cuda_parrots)
+    .done();

mmcv/ops/csrc/parrots/cc_attention_pytorch.h

+#ifndef CC_ATTENTION_PYTORCH_H
+#define CC_ATTENTION_PYTORCH_H
+#include <torch/extension.h>
+using namespace at;
+
+void ca_forward_cuda(const Tensor t, const Tensor f, Tensor weight);
+
+void ca_backward_cuda(const Tensor dw, const Tensor t, const Tensor f,
+                      Tensor dt, Tensor df);
+
+void ca_map_forward_cuda(const Tensor weight, const Tensor g, Tensor out);
+
+void ca_map_backward_cuda(const Tensor dout, const Tensor weight,
+                          const Tensor g, Tensor dw, Tensor dg);
+#endif  // CC_ATTENTION_PYTORCH_H

mmcv/ops/csrc/parrots/corner_pool.cpp

 // Modified from
 // https://github.com/princeton-vl/CornerNet-Lite/tree/master/core/models/py_utils/_cpools/src
-#include "parrots_cpp_helper.hpp"
-
-void bottom_pool_forward_cuda(CudaContext& ctx, const SSElement& attr,
-                              const OperatorBase::in_list_t& ins,
-                              OperatorBase::out_list_t& outs) {}
-
-void bottom_pool_backward_cuda(CudaContext& ctx, const SSElement& attr,
-                               const OperatorBase::in_list_t& ins,
-                               OperatorBase::out_list_t& outs) {}
-
-void top_pool_forward_cuda(CudaContext& ctx, const SSElement& attr,
-                           const OperatorBase::in_list_t& ins,
-                           OperatorBase::out_list_t& outs) {}
-
-void top_pool_backward_cuda(CudaContext& ctx, const SSElement& attr,
-                            const OperatorBase::in_list_t& ins,
-                            OperatorBase::out_list_t& outs) {}
-
-void left_pool_forward_cuda(CudaContext& ctx, const SSElement& attr,
-                            const OperatorBase::in_list_t& ins,
-                            OperatorBase::out_list_t& outs) {}
-
-void left_pool_backward_cuda(CudaContext& ctx, const SSElement& attr,
-                             const OperatorBase::in_list_t& ins,
-                             OperatorBase::out_list_t& outs) {}
-
-void right_pool_forward_cuda(CudaContext& ctx, const SSElement& attr,
-                             const OperatorBase::in_list_t& ins,
-                             OperatorBase::out_list_t& outs) {}
-
-void right_pool_backward_cuda(CudaContext& ctx, const SSElement& attr,
-                              const OperatorBase::in_list_t& ins,
-                              OperatorBase::out_list_t& outs) {}
-
-PARROTS_EXTENSION_REGISTER(bottom_pool_forward)
-    .input(1)
-    .output(1)
-    .apply(bottom_pool_forward_cuda)
-    .done();
-
-PARROTS_EXTENSION_REGISTER(bottom_pool_backward)
-    .input(2)
-    .output(1)
-    .apply(bottom_pool_backward_cuda)
-    .done();
-
-PARROTS_EXTENSION_REGISTER(top_pool_forward)
-    .input(1)
-    .output(1)
-    .apply(top_pool_forward_cuda)
-    .done();
-
-PARROTS_EXTENSION_REGISTER(top_pool_backward)
-    .input(2)
-    .output(1)
-    .apply(top_pool_backward_cuda)
-    .done();
-
-PARROTS_EXTENSION_REGISTER(left_pool_forward)
-    .input(1)
-    .output(1)
-    .apply(left_pool_forward_cuda)
-    .done();
-
-PARROTS_EXTENSION_REGISTER(left_pool_backward)
-    .input(2)
-    .output(1)
-    .apply(left_pool_backward_cuda)
-    .done();
-
-PARROTS_EXTENSION_REGISTER(right_pool_forward)
-    .input(1)
-    .output(1)
-    .apply(right_pool_forward_cuda)
-    .done();
-
-PARROTS_EXTENSION_REGISTER(right_pool_backward)
-    .input(2)
-    .output(1)
-    .apply(right_pool_backward_cuda)
-    .done();
+#include "pytorch_cpp_helper.hpp"
+
+Tensor bottom_pool_forward(Tensor input) {
+  // Initialize output
+  Tensor output = at::zeros_like(input);
+  // Get height
+  int64_t height = input.size(2);
+  output.copy_(input);
+
+  for (int64_t ind = 1; ind < height; ind <<= 1) {
+    Tensor max_temp = at::slice(output, 2, ind, height);
+    Tensor cur_temp = at::slice(output, 2, ind, height).clone();
+    Tensor next_temp = at::slice(output, 2, 0, height - ind).clone();
+    at::max_out(max_temp, cur_temp, next_temp);
+  }
+
+  return output;
+}
+
+Tensor bottom_pool_backward(Tensor input, Tensor grad_output) {
+  auto output = at::zeros_like(input);
+
+  int32_t batch = input.size(0);
+  int32_t channel = input.size(1);
+  int32_t height = input.size(2);
+  int32_t width = input.size(3);
+
+  auto max_val = torch::zeros({batch, channel, width},
+                              at::device(at::kCUDA).dtype(at::kFloat));
+  auto max_ind = torch::zeros({batch, channel, width},
+                              at::device(at::kCUDA).dtype(at::kLong));
+
+  auto input_temp = input.select(2, 0);
+  max_val.copy_(input_temp);
+
+  max_ind.fill_(0);
+
+  auto output_temp = output.select(2, 0);
+  auto grad_output_temp = grad_output.select(2, 0);
+  output_temp.copy_(grad_output_temp);
+
+  auto un_max_ind = max_ind.unsqueeze(2);
+  auto gt_mask = torch::zeros({batch, channel, width},
+                              at::device(at::kCUDA).dtype(at::kBool));
+  auto max_temp = torch::zeros({batch, channel, width},
+                               at::device(at::kCUDA).dtype(at::kFloat));
+  for (int32_t ind = 0; ind < height - 1; ++ind) {
+    input_temp = input.select(2, ind + 1);
+    at::gt_out(gt_mask, input_temp, max_val);
+
+    at::masked_select_out(max_temp, input_temp, gt_mask);
+    max_val.masked_scatter_(gt_mask, max_temp);
+    max_ind.masked_fill_(gt_mask, ind + 1);
+
+    grad_output_temp = grad_output.select(2, ind + 1).unsqueeze(2);
+    output.scatter_add_(2, un_max_ind, grad_output_temp);
+  }
+
+  return output;
+}
+
+Tensor left_pool_forward(Tensor input) {
+  // Initialize output
+  Tensor output = at::zeros_like(input);
+  // Get width
+  int64_t width = input.size(3);
+  output.copy_(input);
+
+  for (int64_t ind = 1; ind < width; ind <<= 1) {
+    Tensor max_temp = at::slice(output, 3, 0, width - ind);
+    Tensor cur_temp = at::slice(output, 3, 0, width - ind).clone();
+    Tensor next_temp = at::slice(output, 3, ind, width).clone();
+    at::max_out(max_temp, cur_temp, next_temp);
+  }
+
+  return output;
+}
+
+Tensor left_pool_backward(Tensor input, Tensor grad_output) {
+  auto output = at::zeros_like(input);
+
+  int32_t batch = input.size(0);
+  int32_t channel = input.size(1);
+  int32_t height = input.size(2);
+  int32_t width = input.size(3);
+
+  auto max_val = torch::zeros({batch, channel, height},
+                              at::device(at::kCUDA).dtype(at::kFloat));
+  auto max_ind = torch::zeros({batch, channel, height},
+                              at::device(at::kCUDA).dtype(at::kLong));
+
+  auto input_temp = input.select(3, width - 1);
+  max_val.copy_(input_temp);
+
+  max_ind.fill_(width - 1);
+
+  auto output_temp = output.select(3, width - 1);
+  auto grad_output_temp = grad_output.select(3, width - 1);
+  output_temp.copy_(grad_output_temp);
+
+  auto un_max_ind = max_ind.unsqueeze(3);
+  auto gt_mask = torch::zeros({batch, channel, height},
+                              at::device(at::kCUDA).dtype(at::kBool));
+  auto max_temp = torch::zeros({batch, channel, height},
+                               at::device(at::kCUDA).dtype(at::kFloat));
+  for (int32_t ind = 1; ind < width; ++ind) {
+    input_temp = input.select(3, width - ind - 1);
+    at::gt_out(gt_mask, input_temp, max_val);
+
+    at::masked_select_out(max_temp, input_temp, gt_mask);
+    max_val.masked_scatter_(gt_mask, max_temp);
+    max_ind.masked_fill_(gt_mask, width - ind - 1);
+
+    grad_output_temp = grad_output.select(3, width - ind - 1).unsqueeze(3);
+    output.scatter_add_(3, un_max_ind, grad_output_temp);
+  }
+
+  return output;
+}
+
+Tensor right_pool_forward(Tensor input) {
+  // Initialize output
+  Tensor output = at::zeros_like(input);
+  // Get width
+  int64_t width = input.size(3);
+  output.copy_(input);
+
+  for (int64_t ind = 1; ind < width; ind <<= 1) {
+    Tensor max_temp = at::slice(output, 3, ind, width);
+    Tensor cur_temp = at::slice(output, 3, ind, width).clone();
+    Tensor next_temp = at::slice(output, 3, 0, width - ind).clone();
+    at::max_out(max_temp, cur_temp, next_temp);
+  }
+
+  return output;
+}
+
+Tensor right_pool_backward(Tensor input, Tensor grad_output) {
+  Tensor output = at::zeros_like(input);
+
+  int32_t batch = input.size(0);
+  int32_t channel = input.size(1);
+  int32_t height = input.size(2);
+  int32_t width = input.size(3);
+
+  auto max_val = torch::zeros({batch, channel, height},
+                              at::device(at::kCUDA).dtype(at::kFloat));
+  auto max_ind = torch::zeros({batch, channel, height},
+                              at::device(at::kCUDA).dtype(at::kLong));
+
+  auto input_temp = input.select(3, 0);
+  max_val.copy_(input_temp);
+
+  max_ind.fill_(0);
+
+  auto output_temp = output.select(3, 0);
+  auto grad_output_temp = grad_output.select(3, 0);
+  output_temp.copy_(grad_output_temp);
+
+  auto un_max_ind = max_ind.unsqueeze(3);
+  auto gt_mask = torch::zeros({batch, channel, height},
+                              at::device(at::kCUDA).dtype(at::kBool));
+  auto max_temp = torch::zeros({batch, channel, height},
+                               at::device(at::kCUDA).dtype(at::kFloat));
+  for (int32_t ind = 0; ind < width - 1; ++ind) {
+    input_temp = input.select(3, ind + 1);
+    at::gt_out(gt_mask, input_temp, max_val);
+
+    at::masked_select_out(max_temp, input_temp, gt_mask);
+    max_val.masked_scatter_(gt_mask, max_temp);
+    max_ind.masked_fill_(gt_mask, ind + 1);
+
+    grad_output_temp = grad_output.select(3, ind + 1).unsqueeze(3);
+    output.scatter_add_(3, un_max_ind, grad_output_temp);
+  }
+
+  return output;
+}
+
+Tensor top_pool_forward(Tensor input) {
+  // Initialize output
+  Tensor output = at::zeros_like(input);
+  // Get height
+  int64_t height = input.size(2);
+  output.copy_(input);
+
+  for (int64_t ind = 1; ind < height; ind <<= 1) {
+    Tensor max_temp = at::slice(output, 2, 0, height - ind);
+    Tensor cur_temp = at::slice(output, 2, 0, height - ind).clone();
+    Tensor next_temp = at::slice(output, 2, ind, height).clone();
+    at::max_out(max_temp, cur_temp, next_temp);
+  }
+
+  return output;
+}
+
+Tensor top_pool_backward(Tensor input, Tensor grad_output) {
+  auto output = at::zeros_like(input);
+
+  int32_t batch = input.size(0);
+  int32_t channel = input.size(1);
+  int32_t height = input.size(2);
+  int32_t width = input.size(3);
+
+  auto max_val = torch::zeros({batch, channel, width},
+                              at::device(at::kCUDA).dtype(at::kFloat));
+  auto max_ind = torch::zeros({batch, channel, width},
+                              at::device(at::kCUDA).dtype(at::kLong));
+
+  auto input_temp = input.select(2, height - 1);
+  max_val.copy_(input_temp);
+
+  max_ind.fill_(height - 1);
+
+  auto output_temp = output.select(2, height - 1);
+  auto grad_output_temp = grad_output.select(2, height - 1);
+  output_temp.copy_(grad_output_temp);
+
+  auto un_max_ind = max_ind.unsqueeze(2);
+  auto gt_mask = torch::zeros({batch, channel, width},
+                              at::device(at::kCUDA).dtype(at::kBool));
+  auto max_temp = torch::zeros({batch, channel, width},
+                               at::device(at::kCUDA).dtype(at::kFloat));
+  for (int32_t ind = 1; ind < height; ++ind) {
+    input_temp = input.select(2, height - ind - 1);
+    at::gt_out(gt_mask, input_temp, max_val);
+
+    at::masked_select_out(max_temp, input_temp, gt_mask);
+    max_val.masked_scatter_(gt_mask, max_temp);
+    max_ind.masked_fill_(gt_mask, height - ind - 1);
+
+    grad_output_temp = grad_output.select(2, height - ind - 1).unsqueeze(2);
+    output.scatter_add_(2, un_max_ind, grad_output_temp);
+  }
+
+  return output;
+}

mmcv/ops/csrc/parrots/corner_pool_parrots.cpp

+#include <parrots/compute/aten.hpp>
+#include <parrots/extension.hpp>
+#include <parrots/foundation/ssattrs.hpp>
+
+#include "corner_pool_pytorch.h"
+
+using namespace parrots;
+
+#ifdef MMCV_WITH_CUDA
+void bottom_pool_forward_parrots(CudaContext& ctx, const SSElement& attr,
+                                 const OperatorBase::in_list_t& ins,
+                                 OperatorBase::out_list_t& outs) {
+  at::Tensor input;
+  input = buildATensor(ctx, ins[0]);
+  auto out = bottom_pool_forward(input);
+  updateDArray(ctx, out, outs[0]);
+}
+
+void bottom_pool_backward_parrots(CudaContext& ctx, const SSElement& attr,
+                                  const OperatorBase::in_list_t& ins,
+                                  OperatorBase::out_list_t& outs) {
+  at::Tensor input, grad_output;
+  input = buildATensor(ctx, ins[0]);
+  grad_output = buildATensor(ctx, ins[1]);
+  auto out = bottom_pool_backward(input, grad_output);
+  updateDArray(ctx, out, outs[0]);
+}
+
+void left_pool_forward_parrots(CudaContext& ctx, const SSElement& attr,
+                               const OperatorBase::in_list_t& ins,
+                               OperatorBase::out_list_t& outs) {
+  at::Tensor input;
+  input = buildATensor(ctx, ins[0]);
+  auto out = left_pool_forward(input);
+  updateDArray(ctx, out, outs[0]);
+}
+
+void left_pool_backward_parrots(CudaContext& ctx, const SSElement& attr,
+                                const OperatorBase::in_list_t& ins,
+                                OperatorBase::out_list_t& outs) {
+  at::Tensor input, grad_output;
+  input = buildATensor(ctx, ins[0]);
+  grad_output = buildATensor(ctx, ins[1]);
+  auto out = left_pool_backward(input, grad_output);
+  updateDArray(ctx, out, outs[0]);
+}
+
+void right_pool_forward_parrots(CudaContext& ctx, const SSElement& attr,
+                                const OperatorBase::in_list_t& ins,
+                                OperatorBase::out_list_t& outs) {
+  at::Tensor input;
+  input = buildATensor(ctx, ins[0]);
+  auto out = right_pool_forward(input);
+  updateDArray(ctx, out, outs[0]);
+}
+
+void right_pool_backward_parrots(CudaContext& ctx, const SSElement& attr,
+                                 const OperatorBase::in_list_t& ins,
+                                 OperatorBase::out_list_t& outs) {
+  at::Tensor input, grad_output;
+  input = buildATensor(ctx, ins[0]);
+  grad_output = buildATensor(ctx, ins[1]);
+  auto out = right_pool_backward(input, grad_output);
+  updateDArray(ctx, out, outs[0]);
+}
+
+void top_pool_forward_parrots(CudaContext& ctx, const SSElement& attr,
+                              const OperatorBase::in_list_t& ins,
+                              OperatorBase::out_list_t& outs) {
+  at::Tensor input;
+  input = buildATensor(ctx, ins[0]);
+  auto out = top_pool_forward(input);
+  updateDArray(ctx, out, outs[0]);
+}
+
+void top_pool_backward_parrots(CudaContext& ctx, const SSElement& attr,
+                               const OperatorBase::in_list_t& ins,
+                               OperatorBase::out_list_t& outs) {
+  at::Tensor input, grad_output;
+  input = buildATensor(ctx, ins[0]);
+  grad_output = buildATensor(ctx, ins[1]);
+  auto out = top_pool_backward(input, grad_output);
+  updateDArray(ctx, out, outs[0]);
+}
+#endif
+
+void bottom_pool_forward_parrots_cpu(HostContext& ctx, const SSElement& attr,
+                                     const OperatorBase::in_list_t& ins,
+                                     OperatorBase::out_list_t& outs) {
+  at::Tensor input;
+  input = buildATensor(ctx, ins[0]);
+  auto out = bottom_pool_forward(input);
+  updateDArray(ctx, out, outs[0]);
+}
+
+void bottom_pool_backward_parrots_cpu(HostContext& ctx, const SSElement& attr,
+                                      const OperatorBase::in_list_t& ins,
+                                      OperatorBase::out_list_t& outs) {
+  at::Tensor input, grad_output;
+  input = buildATensor(ctx, ins[0]);
+  grad_output = buildATensor(ctx, ins[1]);
+  auto out = bottom_pool_backward(input, grad_output);
+  updateDArray(ctx, out, outs[0]);
+}
+
+void left_pool_forward_parrots_cpu(HostContext& ctx, const SSElement& attr,
+                                   const OperatorBase::in_list_t& ins,
+                                   OperatorBase::out_list_t& outs) {
+  at::Tensor input;
+  input = buildATensor(ctx, ins[0]);
+  auto out = left_pool_forward(input);
+  updateDArray(ctx, out, outs[0]);
+}
+
+void left_pool_backward_parrots_cpu(HostContext& ctx, const SSElement& attr,
+                                    const OperatorBase::in_list_t& ins,
+                                    OperatorBase::out_list_t& outs) {
+  at::Tensor input, grad_output;
+  input = buildATensor(ctx, ins[0]);
+  grad_output = buildATensor(ctx, ins[1]);
+  auto out = left_pool_backward(input, grad_output);
+  updateDArray(ctx, out, outs[0]);
+}
+
+void right_pool_forward_parrots_cpu(HostContext& ctx, const SSElement& attr,
+                                    const OperatorBase::in_list_t& ins,
+                                    OperatorBase::out_list_t& outs) {
+  at::Tensor input;
+  input = buildATensor(ctx, ins[0]);
+  auto out = right_pool_forward(input);
+  updateDArray(ctx, out, outs[0]);
+}
+
+void right_pool_backward_parrots_cpu(HostContext& ctx, const SSElement& attr,
+                                     const OperatorBase::in_list_t& ins,
+                                     OperatorBase::out_list_t& outs) {
+  at::Tensor input, grad_output;
+  input = buildATensor(ctx, ins[0]);
+  grad_output = buildATensor(ctx, ins[1]);
+  auto out = right_pool_backward(input, grad_output);
+  updateDArray(ctx, out, outs[0]);
+}
+
+void top_pool_forward_parrots_cpu(HostContext& ctx, const SSElement& attr,
+                                  const OperatorBase::in_list_t& ins,
+                                  OperatorBase::out_list_t& outs) {
+  at::Tensor input;
+  input = buildATensor(ctx, ins[0]);
+  auto out = top_pool_forward(input);
+  updateDArray(ctx, out, outs[0]);
+}
+
+void top_pool_backward_parrots_cpu(HostContext& ctx, const SSElement& attr,
+                                   const OperatorBase::in_list_t& ins,
+                                   OperatorBase::out_list_t& outs) {
+  at::Tensor input, grad_output;
+  input = buildATensor(ctx, ins[0]);
+  grad_output = buildATensor(ctx, ins[1]);
+  auto out = top_pool_backward(input, grad_output);
+  updateDArray(ctx, out, outs[0]);
+}
+
+PARROTS_EXTENSION_REGISTER(bottom_pool_forward)
+    .input(1)
+    .output(1)
+#ifdef MMCV_WITH_CUDA
+    .apply(bottom_pool_forward_parrots)
+#endif
+    .apply(bottom_pool_forward_parrots_cpu)
+    .done();
+
+PARROTS_EXTENSION_REGISTER(bottom_pool_backward)
+    .input(2)
+    .output(1)
+#ifdef MMCV_WITH_CUDA
+    .apply(bottom_pool_backward_parrots)
+#endif
+    .apply(bottom_pool_backward_parrots_cpu)
+    .done();
+
+PARROTS_EXTENSION_REGISTER(top_pool_forward)
+    .input(1)
+    .output(1)
+#ifdef MMCV_WITH_CUDA
+    .apply(top_pool_forward_parrots)
+#endif
+    .apply(top_pool_forward_parrots_cpu)
+    .done();
+
+PARROTS_EXTENSION_REGISTER(top_pool_backward)
+    .input(2)
+    .output(1)
+#ifdef MMCV_WITH_CUDA
+    .apply(top_pool_backward_parrots)
+#endif
+    .apply(top_pool_backward_parrots_cpu)
+    .done();
+
+PARROTS_EXTENSION_REGISTER(left_pool_forward)
+    .input(1)
+    .output(1)
+#ifdef MMCV_WITH_CUDA
+    .apply(left_pool_forward_parrots)
+#endif
+    .apply(left_pool_forward_parrots_cpu)
+    .done();
+
+PARROTS_EXTENSION_REGISTER(left_pool_backward)
+    .input(2)
+    .output(1)
+#ifdef MMCV_WITH_CUDA
+    .apply(left_pool_backward_parrots)
+#endif
+    .apply(left_pool_backward_parrots_cpu)
+    .done();
+
+PARROTS_EXTENSION_REGISTER(right_pool_forward)
+    .input(1)
+    .output(1)
+#ifdef MMCV_WITH_CUDA
+    .apply(right_pool_forward_parrots)
+#endif
+    .apply(right_pool_forward_parrots_cpu)
+    .done();
+
+PARROTS_EXTENSION_REGISTER(right_pool_backward)
+    .input(2)
+    .output(1)
+#ifdef MMCV_WITH_CUDA
+    .apply(right_pool_backward_parrots)
+#endif
+    .apply(right_pool_backward_parrots_cpu)
+    .done();

mmcv/ops/csrc/parrots/corner_pool_pytorch.h

+#ifndef CORNER_POOL_PYTORCH_H
+#define CORNER_POOL_PYTORCH_H
+#include <torch/extension.h>
+
+at::Tensor bottom_pool_forward(at::Tensor input);
+at::Tensor bottom_pool_backward(at::Tensor input, at::Tensor grad_output);
+at::Tensor left_pool_forward(at::Tensor input);
+at::Tensor left_pool_backward(at::Tensor input, at::Tensor grad_output);
+at::Tensor right_pool_forward(at::Tensor input);
+at::Tensor right_pool_backward(at::Tensor input, at::Tensor grad_output);
+at::Tensor top_pool_forward(at::Tensor input);
+at::Tensor top_pool_backward(at::Tensor input, at::Tensor grad_output);
+
+#endif  // CORNER_POOL_PYTORCH_H

mmcv/ops/csrc/parrots/deform_conv.cpp

-// Copyright (c) 2018, SenseTime.
-#include "parrots_cpp_helper.hpp"
+#include "pytorch_cpp_helper.hpp"
 
-void DeformConvForwardCUDAKernelLauncher(
-    const DArrayLite input, const DArrayLite weight, const DArrayLite offset,
-    DArrayLite output, DArrayLite columns, DArrayLite ones, int kW, int kH,
-    int dW, int dH, int padW, int padH, int dilationW, int dilationH, int group,
-    int deformable_group, int im2col_step, CudaContext& ctx,
-    cudaStream_t stream);
+#ifdef MMCV_WITH_CUDA
+void DeformConvForwardCUDAKernelLauncher(Tensor input, Tensor weight,
+                                         Tensor offset, Tensor output,
+                                         Tensor columns, Tensor ones, int kW,
+                                         int kH, int dW, int dH, int padW,
+                                         int padH, int dilationW, int dilationH,
+                                         int group, int deformable_group,
+                                         int im2col_step);
 
 void DeformConvBackwardInputCUDAKernelLauncher(
-    const DArrayLite input, const DArrayLite offset,
-    const DArrayLite gradOutput, DArrayLite gradInput, DArrayLite gradOffset,
-    DArrayLite weight, DArrayLite columns, int kW, int kH, int dW, int dH,
-    int padW, int padH, int dilationW, int dilationH, int group,
-    int deformable_group, int im2col_step, CudaContext& ctx,
-    cudaStream_t stream);
+    Tensor input, Tensor offset, Tensor gradOutput, Tensor gradInput,
+    Tensor gradOffset, Tensor weight, Tensor columns, int kW, int kH, int dW,
+    int dH, int padW, int padH, int dilationW, int dilationH, int group,
+    int deformable_group, int im2col_step);
 
 void DeformConvBackwardParametersCUDAKernelLauncher(
-    const DArrayLite input, const DArrayLite offset,
-    const DArrayLite gradOutput, DArrayLite gradWeight, DArrayLite columns,
-    DArrayLite ones, int kW, int kH, int dW, int dH, int padW, int padH,
-    int dilationW, int dilationH, int group, int deformable_group, float scale,
-    int im2col_step, CudaContext& ctx, cudaStream_t stream);
-
-void deform_conv_forward_cuda(CudaContext& ctx, const SSElement& attr,
-                              const OperatorBase::in_list_t& ins,
-                              OperatorBase::out_list_t& outs) {
-  int kW, kH, dW, dH, padW, padH, dilationW, dilationH, group, deformable_group,
-      im2col_step;
-  SSAttrs(attr)
-      .get<int>("kW", kW)
-      .get<int>("kH", kH)
-      .get<int>("dW", dW)
-      .get<int>("dH", dH)
-      .get<int>("padW", padW)
-      .get<int>("padH", padH)
-      .get<int>("dilationW", dilationW)
-      .get<int>("dilationH", dilationH)
-      .get<int>("group", group)
-      .get<int>("deformable_group", deformable_group)
-      .get<int>("im2col_step", im2col_step)
-      .done();
-
-  const auto input = ins[0];
-  const auto weight = ins[1];
-  const auto offset = ins[2];
-
-  auto output = outs[0];
-  auto columns = outs[1];
-  auto ones = outs[2];
-
-  cudaStream_t stream = getStreamNative<CudaDevice>(ctx.getStream());
+    Tensor input, Tensor offset, Tensor gradOutput, Tensor gradWeight,
+    Tensor columns, Tensor ones, int kW, int kH, int dW, int dH, int padW,
+    int padH, int dilationW, int dilationH, int group, int deformable_group,
+    float scale, int im2col_step);
+
+void deform_conv_forward_cuda(Tensor input, Tensor weight, Tensor offset,
+                              Tensor output, Tensor columns, Tensor ones,
+                              int kW, int kH, int dW, int dH, int padW,
+                              int padH, int dilationW, int dilationH, int group,
+                              int deformable_group, int im2col_step) {
   DeformConvForwardCUDAKernelLauncher(
       input, weight, offset, output, columns, ones, kW, kH, dW, dH, padW, padH,
-      dilationW, dilationH, group, deformable_group, im2col_step, ctx, stream);
+      dilationW, dilationH, group, deformable_group, im2col_step);
 }
 
-void deform_conv_backward_input_cuda(CudaContext& ctx, const SSElement& attr,
-                                     const OperatorBase::in_list_t& ins,
-                                     OperatorBase::out_list_t& outs) {
-  int kW, kH, dW, dH, padW, padH, dilationW, dilationH, group, deformable_group,
-      im2col_step;
-  SSAttrs(attr)
-      .get<int>("kW", kW)
-      .get<int>("kH", kH)
-      .get<int>("dW", dW)
-      .get<int>("dH", dH)
-      .get<int>("padW", padW)
-      .get<int>("padH", padH)
-      .get<int>("dilationW", dilationW)
-      .get<int>("dilationH", dilationH)
-      .get<int>("group", group)
-      .get<int>("deformable_group", deformable_group)
-      .get<int>("im2col_step", im2col_step)
-      .done();
-
-  auto input = ins[0];
-  auto offset = ins[1];
-  auto gradOutput = ins[2];
-
-  auto gradInput = outs[0];
-  auto gradOffset = outs[1];
-  auto weight = outs[2];
-  auto columns = outs[3];
-
-  cudaStream_t stream = getStreamNative<CudaDevice>(ctx.getStream());
+void deform_conv_backward_input_cuda(Tensor input, Tensor offset,
+                                     Tensor gradOutput, Tensor gradInput,
+                                     Tensor gradOffset, Tensor weight,
+                                     Tensor columns, int kW, int kH, int dW,
+                                     int dH, int padW, int padH, int dilationW,
+                                     int dilationH, int group,
+                                     int deformable_group, int im2col_step) {
   DeformConvBackwardInputCUDAKernelLauncher(
       input, offset, gradOutput, gradInput, gradOffset, weight, columns, kW, kH,
       dW, dH, padW, padH, dilationW, dilationH, group, deformable_group,
-      im2col_step, ctx, stream);
+      im2col_step);
 }
 
-void deform_conv_backward_parameters_cuda(CudaContext& ctx,
-                                          const SSElement& attr,
-                                          const OperatorBase::in_list_t& ins,
-                                          OperatorBase::out_list_t& outs) {
-  int kW, kH, dW, dH, padW, padH, dilationW, dilationH, group, deformable_group,
-      im2col_step;
-  float scale;
-  SSAttrs(attr)
-      .get<int>("kW", kW)
-      .get<int>("kH", kH)
-      .get<int>("dW", dW)
-      .get<int>("dH", dH)
-      .get<int>("padW", padW)
-      .get<int>("padH", padH)
-      .get<int>("dilationW", dilationW)
-      .get<int>("dilationH", dilationH)
-      .get<int>("group", group)
-      .get<int>("deformable_group", deformable_group)
-      .get<float>("scale", scale)
-      .get<int>("im2col_step", im2col_step)
-      .done();
-
-  auto input = ins[0];
-  auto offset = ins[1];
-  auto gradOutput = ins[2];
-
-  auto gradWeight = outs[0];
-  auto columns = outs[1];
-  auto ones = outs[2];
-
-  cudaStream_t stream = getStreamNative<CudaDevice>(ctx.getStream());
+void deform_conv_backward_parameters_cuda(
+    Tensor input, Tensor offset, Tensor gradOutput, Tensor gradWeight,
+    Tensor columns, Tensor ones, int kW, int kH, int dW, int dH, int padW,
+    int padH, int dilationW, int dilationH, int group, int deformable_group,
+    float scale, int im2col_step) {
   DeformConvBackwardParametersCUDAKernelLauncher(
       input, offset, gradOutput, gradWeight, columns, ones, kW, kH, dW, dH,
       padW, padH, dilationW, dilationH, group, deformable_group, scale,
-      im2col_step, ctx, stream);
+      im2col_step);
+}
+#endif
+
+void deform_conv_forward(Tensor input, Tensor weight, Tensor offset,
+                         Tensor output, Tensor columns, Tensor ones, int kW,
+                         int kH, int dW, int dH, int padW, int padH,
+                         int dilationW, int dilationH, int group,
+                         int deformable_group, int im2col_step) {
+  if (input.device().is_cuda()) {
+#ifdef MMCV_WITH_CUDA
+    CHECK_CUDA_INPUT(input);
+    CHECK_CUDA_INPUT(offset);
+    CHECK_CUDA_INPUT(weight);
+    CHECK_CUDA_INPUT(output);
+    CHECK_CUDA_INPUT(columns);
+    CHECK_CUDA_INPUT(ones);
+
+    deform_conv_forward_cuda(input, weight, offset, output, columns, ones, kW,
+                             kH, dW, dH, padW, padH, dilationW, dilationH,
+                             group, deformable_group, im2col_step);
+#else
+    AT_ERROR("DeformConv is not compiled with GPU support");
+#endif
+  } else {
+    AT_ERROR("DeformConv is not implemented on CPU");
+  }
 }
 
-PARROTS_EXTENSION_REGISTER(deform_conv_forward)
-    .attr("kW")
-    .attr("kH")
-    .attr("dW")
-    .attr("dH")
-    .attr("padW")
-    .attr("padH")
-    .attr("dilationW")
-    .attr("dilationH")
-    .attr("group")
-    .attr("deformable_group")
-    .attr("im2col_step")
-    .input(3)
-    .output(3)
-    .apply(deform_conv_forward_cuda)
-    .done();
-
-PARROTS_EXTENSION_REGISTER(deform_conv_backward_input)
-    .attr("kW")
-    .attr("kH")
-    .attr("dW")
-    .attr("dH")
-    .attr("padW")
-    .attr("padH")
-    .attr("dilationW")
-    .attr("dilationH")
-    .attr("group")
-    .attr("deformable_group")
-    .attr("im2col_step")
-    .input(3)
-    .output(4)
-    .apply(deform_conv_backward_input_cuda)
-    .done();
+void deform_conv_backward_input(Tensor input, Tensor offset, Tensor gradOutput,
+                                Tensor gradInput, Tensor gradOffset,
+                                Tensor weight, Tensor columns, int kW, int kH,
+                                int dW, int dH, int padW, int padH,
+                                int dilationW, int dilationH, int group,
+                                int deformable_group, int im2col_step) {
+  if (input.device().is_cuda()) {
+#ifdef MMCV_WITH_CUDA
+    CHECK_CUDA_INPUT(input);
+    CHECK_CUDA_INPUT(offset);
+    CHECK_CUDA_INPUT(gradOutput);
+    CHECK_CUDA_INPUT(gradInput);
+    CHECK_CUDA_INPUT(gradOffset);
+    CHECK_CUDA_INPUT(weight);
+    CHECK_CUDA_INPUT(columns);
+
+    deform_conv_backward_input_cuda(input, offset, gradOutput, gradInput,
+                                    gradOffset, weight, columns, kW, kH, dW, dH,
+                                    padW, padH, dilationW, dilationH, group,
+                                    deformable_group, im2col_step);
+#else
+    AT_ERROR("DeformConv is not compiled with GPU support");
+#endif
+  } else {
+    AT_ERROR("DeformConv is not implemented on CPU");
+  }
+}
 
-PARROTS_EXTENSION_REGISTER(deform_conv_backward_parameters)
-    .attr("kW")
-    .attr("kH")
-    .attr("dW")
-    .attr("dH")
-    .attr("padW")
-    .attr("padH")
-    .attr("dilationW")
-    .attr("dilationH")
-    .attr("group")
-    .attr("deformable_group")
-    .attr("scale")
-    .attr("im2col_step")
-    .input(3)
-    .output(3)
-    .apply(deform_conv_backward_parameters_cuda)
-    .done();
+void deform_conv_backward_parameters(Tensor input, Tensor offset,
+                                     Tensor gradOutput, Tensor gradWeight,
+                                     Tensor columns, Tensor ones, int kW,
+                                     int kH, int dW, int dH, int padW, int padH,
+                                     int dilationW, int dilationH, int group,
+                                     int deformable_group, float scale,
+                                     int im2col_step) {
+  if (input.device().is_cuda()) {
+#ifdef MMCV_WITH_CUDA
+    CHECK_CUDA_INPUT(input);
+    CHECK_CUDA_INPUT(offset);
+    CHECK_CUDA_INPUT(gradOutput);
+    CHECK_CUDA_INPUT(gradWeight);
+    CHECK_CUDA_INPUT(columns);
+    CHECK_CUDA_INPUT(ones);
+
+    deform_conv_backward_parameters_cuda(input, offset, gradOutput, gradWeight,
+                                         columns, ones, kW, kH, dW, dH, padW,
+                                         padH, dilationW, dilationH, group,
+                                         deformable_group, scale, im2col_step);
+#else
+    AT_ERROR("DeformConv is not compiled with GPU support");
+#endif
+  } else {
+    AT_ERROR("DeformConv is not implemented on CPU");
+  }
+}

mmcv/ops/csrc/parrots/deform_conv_cuda.cu

 #include "deform_conv_cuda_kernel.cuh"
-#include "parrots_cuda_helper.hpp"
+#include "pytorch_cuda_helper.hpp"
 
-void deformable_im2col(DArrayLite data_im, DArrayLite data_offset,
-                       const int channels, const int height, const int width,
-                       const int ksize_h, const int ksize_w, const int pad_h,
-                       const int pad_w, const int stride_h, const int stride_w,
+void deformable_im2col(Tensor data_im, Tensor data_offset, const int channels,
+                       const int height, const int width, const int ksize_h,
+                       const int ksize_w, const int pad_h, const int pad_w,
+                       const int stride_h, const int stride_w,
                        const int dilation_h, const int dilation_w,
                        const int parallel_imgs, const int deformable_group,
-                       DArrayLite data_col, cudaStream_t stream) {
+                       Tensor data_col) {
   // num_axes should be smaller than block size
   // todo: check parallel_imgs is correctly passed in
   int height_col =
@@ -17,28 +17,31 @@ void deformable_im2col(DArrayLite data_im, DArrayLite data_offset,
   int num_kernels = channels * height_col * width_col * parallel_imgs;
   int channel_per_deformable_group = channels / deformable_group;
 
-  PARROTS_DISPATCH_FLOATING_TYPES_AND_HALF(
-      data_im.elemType().prim(), ([&] {
-        deformable_im2col_gpu_kernel<scalar_t>
-            <<<GET_BLOCKS(num_kernels), THREADS_PER_BLOCK, 0, stream>>>(
-                num_kernels, data_im.ptr<scalar_t>(),
-                data_offset.ptr<scalar_t>(), height, width, ksize_h, ksize_w,
-                pad_h, pad_w, stride_h, stride_w, dilation_h, dilation_w,
-                channel_per_deformable_group, parallel_imgs, channels,
-                deformable_group, height_col, width_col,
-                data_col.ptr<scalar_t>());
-      }));
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+      data_im.scalar_type(), "deformable_im2col_gpu", ([&] {
+        const scalar_t *data_im_ = data_im.data_ptr<scalar_t>();
+        const scalar_t *data_offset_ = data_offset.data_ptr<scalar_t>();
+        scalar_t *data_col_ = data_col.data_ptr<scalar_t>();
 
-  PARROTS_CUDA_CHECK(cudaGetLastError());
+        deformable_im2col_gpu_kernel<<<GET_BLOCKS(num_kernels),
+                                       THREADS_PER_BLOCK, 0,
+                                       at::cuda::getCurrentCUDAStream()>>>(
+            num_kernels, data_im_, data_offset_, height, width, ksize_h,
+            ksize_w, pad_h, pad_w, stride_h, stride_w, dilation_h, dilation_w,
+            channel_per_deformable_group, parallel_imgs, channels,
+            deformable_group, height_col, width_col, data_col_);
+      }));
+  AT_CUDA_CHECK(cudaGetLastError());
 }
 
-void deformable_col2im(DArrayLite data_col, DArrayLite data_offset,
-                       const int channels, const int height, const int width,
-                       const int ksize_h, const int ksize_w, const int pad_h,
-                       const int pad_w, const int stride_h, const int stride_w,
+void deformable_col2im(Tensor data_col, Tensor data_offset, const int channels,
+                       const int height, const int width, const int ksize_h,
+                       const int ksize_w, const int pad_h, const int pad_w,
+                       const int stride_h, const int stride_w,
                        const int dilation_h, const int dilation_w,
                        const int parallel_imgs, const int deformable_group,
-                       DArrayLite grad_im, cudaStream_t stream) {
+                       Tensor grad_im) {
+  // todo: make sure parallel_imgs is passed in correctly
   int height_col =
       (height + 2 * pad_h - (dilation_h * (ksize_h - 1) + 1)) / stride_h + 1;
   int width_col =
@@ -47,27 +50,29 @@ void deformable_col2im(DArrayLite data_col, DArrayLite data_offset,
       channels * ksize_h * ksize_w * height_col * width_col * parallel_imgs;
   int channel_per_deformable_group = channels / deformable_group;
 
-  PARROTS_DISPATCH_FLOATING_TYPES_AND_HALF(
-      data_col.elemType().prim(), ([&] {
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+      data_col.scalar_type(), "deformable_col2im_gpu", ([&] {
+        const scalar_t *data_col_ = data_col.data_ptr<scalar_t>();
+        const scalar_t *data_offset_ = data_offset.data_ptr<scalar_t>();
+        scalar_t *grad_im_ = grad_im.data_ptr<scalar_t>();
+
         deformable_col2im_gpu_kernel<<<GET_BLOCKS(num_kernels),
-                                       THREADS_PER_BLOCK, 0, stream>>>(
-            num_kernels, data_col.ptr<scalar_t>(), data_offset.ptr<scalar_t>(),
-            channels, height, width, ksize_h, ksize_w, pad_h, pad_w, stride_h,
-            stride_w, dilation_h, dilation_w, channel_per_deformable_group,
-            parallel_imgs, deformable_group, height_col, width_col,
-            grad_im.ptr<scalar_t>());
+                                       THREADS_PER_BLOCK, 0,
+                                       at::cuda::getCurrentCUDAStream()>>>(
+            num_kernels, data_col_, data_offset_, channels, height, width,
+            ksize_h, ksize_w, pad_h, pad_w, stride_h, stride_w, dilation_h,
+            dilation_w, channel_per_deformable_group, parallel_imgs,
+            deformable_group, height_col, width_col, grad_im_);
       }));
-
-  PARROTS_CUDA_CHECK(cudaGetLastError());
+  AT_CUDA_CHECK(cudaGetLastError());
 }
 
 void deformable_col2im_coord(
-    DArrayLite data_col, DArrayLite data_im, DArrayLite data_offset,
-    const int channels, const int height, const int width, const int ksize_h,
-    const int ksize_w, const int pad_h, const int pad_w, const int stride_h,
-    const int stride_w, const int dilation_h, const int dilation_w,
-    const int parallel_imgs, const int deformable_group, DArrayLite grad_offset,
-    cudaStream_t stream) {
+    Tensor data_col, Tensor data_im, Tensor data_offset, const int channels,
+    const int height, const int width, const int ksize_h, const int ksize_w,
+    const int pad_h, const int pad_w, const int stride_h, const int stride_w,
+    const int dilation_h, const int dilation_w, const int parallel_imgs,
+    const int deformable_group, Tensor grad_offset) {
   int height_col =
       (height + 2 * pad_h - (dilation_h * (ksize_h - 1) + 1)) / stride_h + 1;
   int width_col =
@@ -77,51 +82,55 @@ void deformable_col2im_coord(
   int channel_per_deformable_group =
       channels * ksize_h * ksize_w / deformable_group;
 
-  PARROTS_DISPATCH_FLOATING_TYPES_AND_HALF(
-      data_col.elemType().prim(), ([&] {
-        deformable_col2im_coord_gpu_kernel<<<GET_BLOCKS(num_kernels),
-                                             THREADS_PER_BLOCK, 0, stream>>>(
-            num_kernels, data_col.ptr<scalar_t>(), data_im.ptr<scalar_t>(),
-            data_offset.ptr<scalar_t>(), channels, height, width, ksize_h,
-            ksize_w, pad_h, pad_w, stride_h, stride_w, dilation_h, dilation_w,
-            channel_per_deformable_group, parallel_imgs,
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+      data_col.scalar_type(), "deformable_col2im_coord_gpu", ([&] {
+        const scalar_t *data_col_ = data_col.data_ptr<scalar_t>();
+        const scalar_t *data_im_ = data_im.data_ptr<scalar_t>();
+        const scalar_t *data_offset_ = data_offset.data_ptr<scalar_t>();
+        scalar_t *grad_offset_ = grad_offset.data_ptr<scalar_t>();
+
+        deformable_col2im_coord_gpu_kernel<<<
+            GET_BLOCKS(num_kernels), THREADS_PER_BLOCK, 0,
+            at::cuda::getCurrentCUDAStream()>>>(
+            num_kernels, data_col_, data_im_, data_offset_, channels, height,
+            width, ksize_h, ksize_w, pad_h, pad_w, stride_h, stride_w,
+            dilation_h, dilation_w, channel_per_deformable_group, parallel_imgs,
             2 * ksize_h * ksize_w * deformable_group, deformable_group,
-            height_col, width_col, grad_offset.ptr<scalar_t>());
+            height_col, width_col, grad_offset_);
       }));
-
-  PARROTS_CUDA_CHECK(cudaGetLastError());
+  AT_CUDA_CHECK(cudaGetLastError());
 }
 
-void deform_conv_shape_check(DArrayLite input, DArrayLite offset,
-                             DArrayLite* gradOutput, DArrayLite weight, int kH,
-                             int kW, int dH, int dW, int padH, int padW,
-                             int dilationH, int dilationW, int group,
-                             int deformable_group) {
-  PARROTS_CHECKARGS(weight.ndims() == 4)
-      << "4D weight tensor (nOutputPlane,nInputPlane,kH,kW) expected, but got: "
-      << weight.ndims();
+void deform_conv_shape_check(Tensor input, Tensor offset, Tensor *gradOutput,
+                             Tensor weight, int kH, int kW, int dH, int dW,
+                             int padH, int padW, int dilationH, int dilationW,
+                             int group, int deformable_group) {
+  TORCH_CHECK(
+      weight.ndimension() == 4,
+      "4D weight tensor (nOutputPlane,nInputPlane,kH,kW) expected, but got: %s",
+      weight.ndimension());
 
-  PARROTS_CHECKARGS(weight.isContiguous())
-      << "weight tensor has to be contiguous";
+  TORCH_CHECK(weight.is_contiguous(), "weight tensor has to be contiguous");
 
-  PARROTS_CHECKARGS(kW > 0 && kH > 0)
-      << "kernel size should be greater than zero, but got kH: " << kH
-      << " kW: " << kW;
+  TORCH_CHECK(kW > 0 && kH > 0,
+              "kernel size should be greater than zero, but got kH: %d kW: %d",
+              kH, kW);
 
-  PARROTS_CHECKARGS(weight.dim(2) == kH && weight.dim(3) == kW)
-      << "kernel size should be consistent with weight, but got kH: " << kH
-      << " kW: " << kW << " weight.dim(2): " << weight.dim(2)
-      << ", weight.dim(3): " << weight.dim(3);
+  TORCH_CHECK((weight.size(2) == kH && weight.size(3) == kW),
+              "kernel size should be consistent with weight, ",
+              "but got kH: %d kW: %d weight.size(2): %d, weight.size(3): %d",
+              kH, kW, weight.size(2), weight.size(3));
 
-  PARROTS_CHECKARGS(dW > 0 && dH > 0)
-      << "stride should be greater than zero, but got dH: " << dH
-      << " dW: " << dW;
+  TORCH_CHECK(dW > 0 && dH > 0,
+              "stride should be greater than zero, but got dH: %d dW: %d", dH,
+              dW);
 
-  PARROTS_CHECKARGS(dilationW > 0 && dilationH > 0)
-      << "dilation should be greater than 0, but got dilationH: " << dilationH
-      << " dilationW: " << dilationW;
+  TORCH_CHECK(
+      dilationW > 0 && dilationH > 0,
+      "dilation should be greater than 0, but got dilationH: %d dilationW: %d",
+      dilationH, dilationW);
 
-  int ndim = input.ndims();
+  int ndim = input.ndimension();
   int dimf = 0;
   int dimh = 1;
   int dimw = 2;
@@ -132,62 +141,67 @@ void deform_conv_shape_check(DArrayLite input, DArrayLite offset,
     dimw++;
   }
 
-  PARROTS_CHECKARGS(ndim == 3 || ndim == 4)
-      << "3D or 4D input tensor expected but got: " << ndim;
+  TORCH_CHECK(ndim == 3 || ndim == 4,
+              "3D or 4D input tensor expected but got: %s", ndim);
 
-  size_t nInputPlane = weight.dim(1) * group;
-  size_t inputHeight = input.dim(dimh);
-  size_t inputWidth = input.dim(dimw);
-  size_t nOutputPlane = weight.dim(0);
-  size_t outputHeight =
+  long nInputPlane = weight.size(1) * group;
+  long inputHeight = input.size(dimh);
+  long inputWidth = input.size(dimw);
+  long nOutputPlane = weight.size(0);
+  long outputHeight =
       (inputHeight + 2 * padH - (dilationH * (kH - 1) + 1)) / dH + 1;
-  size_t outputWidth =
+  long outputWidth =
       (inputWidth + 2 * padW - (dilationW * (kW - 1) + 1)) / dW + 1;
 
-  PARROTS_CHECKARGS(nInputPlane % deformable_group == 0)
-      << "input channels must divide deformable group size";
+  TORCH_CHECK(nInputPlane % deformable_group == 0,
+              "input channels must divide deformable group size");
 
-  PARROTS_CHECKARGS(outputWidth >= 1 || outputHeight >= 1)
-      << "Given input size: (" << nInputPlane << " x " << inputHeight << " x "
-      << inputWidth << "). Calculated output size: (" << nOutputPlane << " x "
-      << outputHeight << " x " << outputWidth << "). Output size is too small";
+  if (outputWidth < 1 || outputHeight < 1)
+    AT_ERROR(
+        "Given input size: (%ld x %ld x %ld). "
+        "Calculated output size: (%ld x %ld x %ld). Output size is too small",
+        nInputPlane, inputHeight, inputWidth, nOutputPlane, outputHeight,
+        outputWidth);
 
-  PARROTS_CHECKARGS(input.dim(1) == nInputPlane)
-      << "invalid number of input planes, expected: " << nInputPlane
-      << ", but got: " << input.dim(1);
+  TORCH_CHECK(input.size(1) == nInputPlane,
+              "invalid number of input planes, expected: %d, but got: %d",
+              nInputPlane, input.size(1));
 
-  PARROTS_CHECKARGS(inputHeight >= kH && inputWidth >= kW)
-      << "input image is smaller than kernel";
+  TORCH_CHECK((inputHeight >= kH && inputWidth >= kW),
+              "input image is smaller than kernel");
 
-  PARROTS_CHECKARGS(offset.dim(2) == outputHeight &&
-                    offset.dim(3) == outputWidth)
-      << "invalid spatial dim of offset, expected height: " << outputHeight
-      << " width: " << outputWidth << ", but got height: " << offset.dim(2)
-      << " width: " << offset.dim(3);
+  TORCH_CHECK(
+      (offset.size(2) == outputHeight && offset.size(3) == outputWidth),
+      "invalid spatial size of offset, expected height: %d width: %d, but "
+      "got height: %d width: %d",
+      outputHeight, outputWidth, offset.size(2), offset.size(3));
 
-  PARROTS_CHECKARGS(offset.dim(1) == deformable_group * 2 * kH * kW)
-      << "invalid number of channels of offset";
+  TORCH_CHECK((offset.size(1) == deformable_group * 2 * kH * kW),
+              "invalid number of channels of offset");
 
   if (gradOutput != NULL) {
-    PARROTS_CHECKARGS(gradOutput->dim(dimf) == nOutputPlane)
-        << "invalid number of gradOutput planes, expected: " << nOutputPlane
-        << ", but got: " << gradOutput->dim(dimf);
-
-    PARROTS_CHECKARGS(gradOutput->dim(dimh) == outputHeight &&
-                      gradOutput->dim(dimw) == outputWidth)
-        << "invalid dim of gradOutput, expected height: " << outputHeight
-        << " width: " << outputWidth
-        << " , but got height: " << gradOutput->dim(dimh)
-        << " width: " << gradOutput->dim(dimw);
+    TORCH_CHECK(
+        gradOutput->size(dimf) == nOutputPlane,
+        "invalid number of gradOutput planes, expected: %d, but got: %d",
+        nOutputPlane, gradOutput->size(dimf));
+
+    TORCH_CHECK(
+        (gradOutput->size(dimh) == outputHeight &&
+         gradOutput->size(dimw) == outputWidth),
+        "invalid size of gradOutput, expected height: %d width: %d , but "
+        "got height: %d width: %d",
+        outputHeight, outputWidth, gradOutput->size(dimh),
+        gradOutput->size(dimw));
   }
 }
 
-void DeformConvForwardCUDAKernelLauncher(
-    DArrayLite input, DArrayLite weight, DArrayLite offset, DArrayLite output,
-    DArrayLite columns, DArrayLite ones, int kW, int kH, int dW, int dH,
-    int padW, int padH, int dilationW, int dilationH, int group,
-    int deformable_group, int im2col_step, CudaContext& ctx,
-    cudaStream_t stream) {
+void DeformConvForwardCUDAKernelLauncher(Tensor input, Tensor weight,
+                                         Tensor offset, Tensor output,
+                                         Tensor columns, Tensor ones, int kW,
+                                         int kH, int dW, int dH, int padW,
+                                         int padH, int dilationW, int dilationH,
+                                         int group, int deformable_group,
+                                         int im2col_step) {
   // todo: resize columns to include im2col: done
   // todo: add im2col_step as input
   // todo: add new output buffer and transpose it to output (or directly
@@ -196,45 +210,41 @@ void DeformConvForwardCUDAKernelLauncher(
 
   deform_conv_shape_check(input, offset, NULL, weight, kH, kW, dH, dW, padH,
                           padW, dilationH, dilationW, group, deformable_group);
+  at::DeviceGuard guard(input.device());
 
   int batch = 1;
-  if (input.ndims() == 3) {
+  if (input.ndimension() == 3) {
     // Force batch
     batch = 0;
-    input = input.view({1, input.dim(0), input.dim(1), input.dim(2)});
-    offset = offset.view({1, offset.dim(0), offset.dim(1), offset.dim(2)});
+    input.unsqueeze_(0);
+    offset.unsqueeze_(0);
   }
 
   // todo: assert batchsize dividable by im2col_step
 
-  size_t batchSize = input.dim(0);
-  size_t nInputPlane = input.dim(1);
-  size_t inputHeight = input.dim(2);
-  size_t inputWidth = input.dim(3);
+  long batchSize = input.size(0);
+  long nInputPlane = input.size(1);
+  long inputHeight = input.size(2);
+  long inputWidth = input.size(3);
 
-  size_t nOutputPlane = weight.dim(0);
+  long nOutputPlane = weight.size(0);
 
-  size_t outputWidth =
+  long outputWidth =
       (inputWidth + 2 * padW - (dilationW * (kW - 1) + 1)) / dW + 1;
-  size_t outputHeight =
+  long outputHeight =
       (inputHeight + 2 * padH - (dilationH * (kH - 1) + 1)) / dH + 1;
 
-  PARROTS_CHECKARGS(offset.dim(0) == batchSize)
-      << "invalid batch size of offset";
+  TORCH_CHECK((offset.size(0) == batchSize), "invalid batch size of offset");
 
   output = output.view({batchSize / im2col_step, im2col_step, nOutputPlane,
                         outputHeight, outputWidth});
+  columns = at::zeros(
+      {nInputPlane * kW * kH, im2col_step * outputHeight * outputWidth},
+      input.options());
 
-  columns = ctx.createDArrayLite(
-      input.elemType(), DArrayShape(nInputPlane * kW * kH,
-                                    im2col_step * outputHeight * outputWidth));
-  columns.setZeros(ctx.getStream());
-
-  if (ones.ndims() != 2 ||
-      ones.dim(0) * ones.dim(1) < outputHeight * outputWidth) {
-    ones = ctx.createDArrayLite(input.elemType(),
-                                DArrayShape(outputHeight, outputWidth));
-    fill(ctx, ones, *toScalar(1));
+  if (ones.ndimension() != 2 ||
+      ones.size(0) * ones.size(1) < outputHeight * outputWidth) {
+    ones = at::ones({outputHeight, outputWidth}, input.options());
   }
 
   input = input.view({batchSize / im2col_step, im2col_step, nInputPlane,
@@ -243,45 +253,41 @@ void DeformConvForwardCUDAKernelLauncher(
       offset.view({batchSize / im2col_step, im2col_step,
                    deformable_group * 2 * kH * kW, outputHeight, outputWidth});
 
-  auto output_buffer = ctx.createDArrayLite(
-      input.elemType(), DArrayShape(batchSize / im2col_step, nOutputPlane,
-                                    im2col_step * outputHeight, outputWidth));
-  output_buffer.setZeros(ctx.getStream());
+  Tensor output_buffer = at::zeros({batchSize / im2col_step, nOutputPlane,
+                                    im2col_step * outputHeight, outputWidth},
+                                   output.options());
+
   output_buffer = output_buffer.view(
-      {output_buffer.dim(0), group, output_buffer.dim(1) / group,
-       output_buffer.dim(2) * output_buffer.dim(3)});
+      {output_buffer.size(0), group, output_buffer.size(1) / group,
+       output_buffer.size(2), output_buffer.size(3)});
 
-  for (size_t elt = 0; elt < batchSize / im2col_step; elt++) {
+  for (int elt = 0; elt < batchSize / im2col_step; elt++) {
     deformable_im2col(input[elt], offset[elt], nInputPlane, inputHeight,
                       inputWidth, kH, kW, padH, padW, dH, dW, dilationH,
-                      dilationW, im2col_step, deformable_group, columns,
-                      stream);
-
-    columns = columns.view({group, columns.dim(0) / group, columns.dim(1)});
-    weight = weight.view(
-        {group, nOutputPlane / group, nInputPlane / group * kH * kW});
-
-    for (size_t g = 0; g < group; g++) {
-      auto output_g = output_buffer[elt][g];
-      auto weight_g = weight[g];
-      auto columns_g = columns[g];
-      gemm(ctx, 1, false, weight_g, false, columns_g, 1, output_g);
+                      dilationW, im2col_step, deformable_group, columns);
+
+    columns = columns.view({group, columns.size(0) / group, columns.size(1)});
+    weight = weight.view({group, weight.size(0) / group, weight.size(1),
+                          weight.size(2), weight.size(3)});
+
+    for (int g = 0; g < group; g++) {
+      output_buffer[elt][g] = output_buffer[elt][g]
+                                  .flatten(1)
+                                  .addmm_(weight[g].flatten(1), columns[g])
+                                  .view_as(output_buffer[elt][g]);
     }
-    columns = columns.view({columns.dim(0) * columns.dim(1), columns.dim(2)});
-    weight = weight.view({nOutputPlane, nInputPlane, kH, kW});
+    columns =
+        columns.view({columns.size(0) * columns.size(1), columns.size(2)});
   }
 
   output_buffer = output_buffer.view(
-      {output_buffer.dim(0), output_buffer.dim(1) * output_buffer.dim(2),
-       output_buffer.dim(3)});
+      {output_buffer.size(0), output_buffer.size(1) * output_buffer.size(2),
+       output_buffer.size(3), output_buffer.size(4)});
 
   output_buffer = output_buffer.view({batchSize / im2col_step, nOutputPlane,
                                       im2col_step, outputHeight, outputWidth});
-  output_buffer = transpose(ctx, output_buffer, 1, 2);
-  if (!output_buffer.isContiguous()) {
-    output_buffer = ctx.cloneDArrayLite(output_buffer);
-  }
-  copy(ctx, output, output_buffer);
+  output_buffer.transpose_(1, 2);
+  output.copy_(output_buffer);
   output = output.view({batchSize, nOutputPlane, outputHeight, outputWidth});
 
   input = input.view({batchSize, nInputPlane, inputHeight, inputWidth});
@@ -291,58 +297,53 @@ void DeformConvForwardCUDAKernelLauncher(
   if (batch == 0) {
     output = output.view({nOutputPlane, outputHeight, outputWidth});
     input = input.view({nInputPlane, inputHeight, inputWidth});
-    offset = offset.view({offset.dim(1), offset.dim(2), offset.dim(3)});
+    offset = offset.view({offset.size(1), offset.size(2), offset.size(3)});
   }
 }
 
 void DeformConvBackwardInputCUDAKernelLauncher(
-    DArrayLite input, DArrayLite offset, DArrayLite gradOutput,
-    DArrayLite gradInput, DArrayLite gradOffset, DArrayLite weight,
-    DArrayLite columns, int kW, int kH, int dW, int dH, int padW, int padH,
-    int dilationW, int dilationH, int group, int deformable_group,
-    int im2col_step, CudaContext& ctx, cudaStream_t stream) {
+    Tensor input, Tensor offset, Tensor gradOutput, Tensor gradInput,
+    Tensor gradOffset, Tensor weight, Tensor columns, int kW, int kH, int dW,
+    int dH, int padW, int padH, int dilationW, int dilationH, int group,
+    int deformable_group, int im2col_step) {
   deform_conv_shape_check(input, offset, &gradOutput, weight, kH, kW, dH, dW,
                           padH, padW, dilationH, dilationW, group,
                           deformable_group);
+  at::DeviceGuard guard(input.device());
 
   int batch = 1;
 
-  if (input.ndims() == 3) {
+  if (input.ndimension() == 3) {
     // Force batch
     batch = 0;
-    input = input.view({1, input.dim(0), input.dim(1), input.dim(2)});
-    offset = offset.view({1, offset.dim(0), offset.dim(1), offset.dim(2)});
+    input = input.view({1, input.size(0), input.size(1), input.size(2)});
+    offset = offset.view({1, offset.size(0), offset.size(1), offset.size(2)});
     gradOutput = gradOutput.view(
-        {1, gradOutput.dim(0), gradOutput.dim(1), gradOutput.dim(2)});
+        {1, gradOutput.size(0), gradOutput.size(1), gradOutput.size(2)});
   }
 
-  size_t batchSize = input.dim(0);
-  size_t nInputPlane = input.dim(1);
-  size_t inputHeight = input.dim(2);
-  size_t inputWidth = input.dim(3);
+  long batchSize = input.size(0);
+  long nInputPlane = input.size(1);
+  long inputHeight = input.size(2);
+  long inputWidth = input.size(3);
 
-  size_t nOutputPlane = weight.dim(0);
+  long nOutputPlane = weight.size(0);
 
-  size_t outputWidth =
+  long outputWidth =
       (inputWidth + 2 * padW - (dilationW * (kW - 1) + 1)) / dW + 1;
-  size_t outputHeight =
+  long outputHeight =
       (inputHeight + 2 * padH - (dilationH * (kH - 1) + 1)) / dH + 1;
 
-  PARROTS_CHECKARGS(offset.dim(0) == batchSize)
-      << "invalid batch size of offset";
+  TORCH_CHECK((offset.size(0) == batchSize), 3, "invalid batch size of offset");
   gradInput = gradInput.view({batchSize, nInputPlane, inputHeight, inputWidth});
-  columns = ctx.createDArrayLite(
-      input.elemType(), DArrayShape(nInputPlane * kW * kH,
-                                    im2col_step * outputHeight * outputWidth));
-  columns.setZeros(ctx.getStream());
+  columns = at::zeros(
+      {nInputPlane * kW * kH, im2col_step * outputHeight * outputWidth},
+      input.options());
 
   // change order of grad output
   gradOutput = gradOutput.view({batchSize / im2col_step, im2col_step,
                                 nOutputPlane, outputHeight, outputWidth});
-  gradOutput = transpose(ctx, gradOutput, 1, 2);
-  if (!gradOutput.isContiguous()) {
-    gradOutput = ctx.cloneDArrayLite(gradOutput);
-  }
+  gradOutput.transpose_(1, 2);
 
   gradInput = gradInput.view({batchSize / im2col_step, im2col_step, nInputPlane,
                               inputHeight, inputWidth});
@@ -355,41 +356,37 @@ void DeformConvBackwardInputCUDAKernelLauncher(
       offset.view({batchSize / im2col_step, im2col_step,
                    deformable_group * 2 * kH * kW, outputHeight, outputWidth});
 
-  for (size_t elt = 0; elt < batchSize / im2col_step; elt++) {
+  for (int elt = 0; elt < batchSize / im2col_step; elt++) {
     // divide into groups
-    columns = columns.view({group, columns.dim(0) / group, columns.dim(1)});
-    weight = weight.view({group, weight.dim(0) / group,
-                          weight.dim(1) * weight.dim(2) * weight.dim(3)});
+    columns = columns.view({group, columns.size(0) / group, columns.size(1)});
+    weight = weight.view({group, weight.size(0) / group, weight.size(1),
+                          weight.size(2), weight.size(3)});
     gradOutput = gradOutput.view(
-        {gradOutput.dim(0), group, gradOutput.dim(1) / group,
-         gradOutput.dim(2) * gradOutput.dim(3) * gradOutput.dim(4)});
+        {gradOutput.size(0), group, gradOutput.size(1) / group,
+         gradOutput.size(2), gradOutput.size(3), gradOutput.size(4)});
 
-    for (size_t g = 0; g < group; g++) {
-      auto columns_g = columns[g];
-      gemm(ctx, 1, true, weight[g], false, gradOutput[elt][g], 0, columns_g);
+    for (int g = 0; g < group; g++) {
+      columns[g] = columns[g].addmm_(weight[g].flatten(1).transpose(0, 1),
+                                     gradOutput[elt][g].flatten(1), 0.0f, 1.0f);
     }
 
-    columns = columns.view({columns.dim(0) * columns.dim(1), columns.dim(2)});
-    gradOutput = gradOutput.view({gradOutput.dim(0),
-                                  gradOutput.dim(1) * gradOutput.dim(2),
-                                  im2col_step, outputHeight, outputWidth});
-    weight = weight.view({nOutputPlane, nInputPlane, kH, kW});
+    columns =
+        columns.view({columns.size(0) * columns.size(1), columns.size(2)});
+    gradOutput = gradOutput.view(
+        {gradOutput.size(0), gradOutput.size(1) * gradOutput.size(2),
+         gradOutput.size(3), gradOutput.size(4), gradOutput.size(5)});
 
     deformable_col2im_coord(columns, input[elt], offset[elt], nInputPlane,
                             inputHeight, inputWidth, kH, kW, padH, padW, dH, dW,
                             dilationH, dilationW, im2col_step, deformable_group,
-                            gradOffset[elt], stream);
+                            gradOffset[elt]);
 
     deformable_col2im(columns, offset[elt], nInputPlane, inputHeight,
                       inputWidth, kH, kW, padH, padW, dH, dW, dilationH,
-                      dilationW, im2col_step, deformable_group, gradInput[elt],
-                      stream);
+                      dilationW, im2col_step, deformable_group, gradInput[elt]);
   }
 
-  gradOutput = transpose(ctx, gradOutput, 1, 2);
-  if (!gradOutput.isContiguous()) {
-    gradOutput = ctx.cloneDArrayLite(gradOutput);
-  }
+  gradOutput.transpose_(1, 2);
   gradOutput =
       gradOutput.view({batchSize, nOutputPlane, outputHeight, outputWidth});
 
@@ -404,17 +401,17 @@ void DeformConvBackwardInputCUDAKernelLauncher(
     gradOutput = gradOutput.view({nOutputPlane, outputHeight, outputWidth});
     input = input.view({nInputPlane, inputHeight, inputWidth});
     gradInput = gradInput.view({nInputPlane, inputHeight, inputWidth});
-    offset = offset.view({offset.dim(1), offset.dim(2), offset.dim(3)});
-    gradOffset = gradOffset.view({offset.dim(1), offset.dim(2), offset.dim(3)});
+    offset = offset.view({offset.size(1), offset.size(2), offset.size(3)});
+    gradOffset =
+        gradOffset.view({offset.size(1), offset.size(2), offset.size(3)});
   }
 }
 
 void DeformConvBackwardParametersCUDAKernelLauncher(
-    DArrayLite input, DArrayLite offset, DArrayLite gradOutput,
-    DArrayLite gradWeight, DArrayLite columns, DArrayLite ones, int kW, int kH,
-    int dW, int dH, int padW, int padH, int dilationW, int dilationH, int group,
-    int deformable_group, float scale, int im2col_step, CudaContext& ctx,
-    cudaStream_t stream) {
+    Tensor input, Tensor offset, Tensor gradOutput, Tensor gradWeight,
+    Tensor columns, Tensor ones, int kW, int kH, int dW, int dH, int padW,
+    int padH, int dilationW, int dilationH, int group, int deformable_group,
+    float scale, int im2col_step) {
   // todo: transpose and reshape outGrad
   // todo: reshape columns
   // todo: add im2col_step as input
@@ -422,53 +419,52 @@ void DeformConvBackwardParametersCUDAKernelLauncher(
   deform_conv_shape_check(input, offset, &gradOutput, gradWeight, kH, kW, dH,
                           dW, padH, padW, dilationH, dilationW, group,
                           deformable_group);
+  at::DeviceGuard guard(input.device());
 
   int batch = 1;
 
-  if (input.ndims() == 3) {
+  if (input.ndimension() == 3) {
     // Force batch
     batch = 0;
-    input = input.view({1, input.dim(0), input.dim(1), input.dim(2)});
+    input = input.view(
+        at::IntList({1, input.size(0), input.size(1), input.size(2)}));
     gradOutput = gradOutput.view(
-        {1, gradOutput.dim(0), gradOutput.dim(1), gradOutput.dim(2)});
+        {1, gradOutput.size(0), gradOutput.size(1), gradOutput.size(2)});
   }
 
-  size_t batchSize = input.dim(0);
-  size_t nInputPlane = input.dim(1);
-  size_t inputHeight = input.dim(2);
-  size_t inputWidth = input.dim(3);
+  long batchSize = input.size(0);
+  long nInputPlane = input.size(1);
+  long inputHeight = input.size(2);
+  long inputWidth = input.size(3);
 
-  size_t nOutputPlane = gradWeight.dim(0);
+  long nOutputPlane = gradWeight.size(0);
 
-  size_t outputWidth =
+  long outputWidth =
       (inputWidth + 2 * padW - (dilationW * (kW - 1) + 1)) / dW + 1;
-  size_t outputHeight =
+  long outputHeight =
       (inputHeight + 2 * padH - (dilationH * (kH - 1) + 1)) / dH + 1;
 
-  PARROTS_CHECKARGS(offset.dim(0) == batchSize)
-      << "invalid batch size of offset";
+  TORCH_CHECK((offset.size(0) == batchSize), "invalid batch size of offset");
 
-  columns = ctx.createDArrayLite(
-      input.elemType(), DArrayShape(nInputPlane * kW * kH,
-                                    im2col_step * outputHeight * outputWidth));
-  columns.setZeros(ctx.getStream());
+  columns = at::zeros(
+      {nInputPlane * kW * kH, im2col_step * outputHeight * outputWidth},
+      input.options());
 
   gradOutput = gradOutput.view({batchSize / im2col_step, im2col_step,
                                 nOutputPlane, outputHeight, outputWidth});
-  gradOutput = transpose(ctx, gradOutput, 1, 2);
-  if (!gradOutput.isContiguous()) {
-    gradOutput = ctx.cloneDArrayLite(gradOutput);
-  }
+  gradOutput.transpose_(1, 2);
 
-  auto gradOutputBuffer = ctx.cloneDArrayLite(gradOutput);
+  Tensor gradOutputBuffer = at::zeros_like(gradOutput);
+  gradOutputBuffer =
+      gradOutputBuffer.view({batchSize / im2col_step, nOutputPlane, im2col_step,
+                             outputHeight, outputWidth});
+  gradOutputBuffer = gradOutputBuffer.contiguous();
+  gradOutputBuffer.copy_(gradOutput);
   gradOutputBuffer =
       gradOutputBuffer.view({batchSize / im2col_step, nOutputPlane,
                              im2col_step * outputHeight, outputWidth});
 
-  gradOutput = transpose(ctx, gradOutput, 1, 2);
-  if (!gradOutput.isContiguous()) {
-    gradOutput = ctx.cloneDArrayLite(gradOutput);
-  }
+  gradOutput.transpose_(1, 2);
   gradOutput =
       gradOutput.view({batchSize, nOutputPlane, outputHeight, outputWidth});
 
@@ -478,33 +474,36 @@ void DeformConvBackwardParametersCUDAKernelLauncher(
       offset.view({batchSize / im2col_step, im2col_step,
                    deformable_group * 2 * kH * kW, outputHeight, outputWidth});
 
-  for (size_t elt = 0; elt < batchSize / im2col_step; elt++) {
+  for (int elt = 0; elt < batchSize / im2col_step; elt++) {
     deformable_im2col(input[elt], offset[elt], nInputPlane, inputHeight,
                       inputWidth, kH, kW, padH, padW, dH, dW, dilationH,
-                      dilationW, im2col_step, deformable_group, columns,
-                      stream);
+                      dilationW, im2col_step, deformable_group, columns);
 
     // divide into group
     gradOutputBuffer = gradOutputBuffer.view(
-        {gradOutputBuffer.dim(0), group, gradOutputBuffer.dim(1) / group,
-         gradOutputBuffer.dim(2) * gradOutputBuffer.dim(3)});
-    columns = columns.view({group, columns.dim(0) / group, columns.dim(1)});
-    gradWeight = gradWeight.view(
-        {group, gradWeight.dim(0) / group,
-         gradWeight.dim(1) * gradWeight.dim(2) * gradWeight.dim(3)});
+        {gradOutputBuffer.size(0), group, gradOutputBuffer.size(1) / group,
+         gradOutputBuffer.size(2), gradOutputBuffer.size(3)});
+    columns = columns.view({group, columns.size(0) / group, columns.size(1)});
+    gradWeight =
+        gradWeight.view({group, gradWeight.size(0) / group, gradWeight.size(1),
+                         gradWeight.size(2), gradWeight.size(3)});
 
     for (int g = 0; g < group; g++) {
-      auto gradWeight_g = gradWeight[g];
-      gemm(ctx, scale, false, gradOutputBuffer[elt][g], true, columns[g], 1,
-           gradWeight_g);
+      gradWeight[g] = gradWeight[g]
+                          .flatten(1)
+                          .addmm_(gradOutputBuffer[elt][g].flatten(1),
+                                  columns[g].transpose(1, 0), 1.0, scale)
+                          .view_as(gradWeight[g]);
     }
     gradOutputBuffer = gradOutputBuffer.view(
-        {gradOutputBuffer.dim(0),
-         gradOutputBuffer.dim(1) * gradOutputBuffer.dim(2),
-         im2col_step * outputHeight, outputWidth});
-    columns = columns.view({columns.dim(0) * columns.dim(1), columns.dim(2)});
-    gradWeight = gradWeight.view(
-        {gradWeight.dim(0) * gradWeight.dim(1), nInputPlane / group, kH, kW});
+        {gradOutputBuffer.size(0),
+         gradOutputBuffer.size(1) * gradOutputBuffer.size(2),
+         gradOutputBuffer.size(3), gradOutputBuffer.size(4)});
+    columns =
+        columns.view({columns.size(0) * columns.size(1), columns.size(2)});
+    gradWeight = gradWeight.view({gradWeight.size(0) * gradWeight.size(1),
+                                  gradWeight.size(2), gradWeight.size(3),
+                                  gradWeight.size(4)});
   }
 
   input = input.view({batchSize, nInputPlane, inputHeight, inputWidth});

mmcv/ops/csrc/parrots/deform_conv_parrots.cpp

+#include <parrots/compute/aten.hpp>
+#include <parrots/extension.hpp>
+#include <parrots/foundation/ssattrs.hpp>
+
+#include "deform_conv_pytorch.h"
+
+using namespace parrots;
+
+/*void deform_conv_forward_cuda(Tensor input, Tensor weight, Tensor offset,
+ *                              Tensor output, Tensor columns, Tensor ones,
+ *                              int kW, int kH, int dW, int dH, int padW,
+ *                              int padH, int dilationW, int dilationH, int
+ * group, int deformable_group, int im2col_step);
+ */
+void deform_conv_forward_cuda_parrots(CudaContext& ctx, const SSElement& attr,
+                                      const OperatorBase::in_list_t& ins,
+                                      OperatorBase::out_list_t& outs) {
+  int kW, kH, dW, dH, padW, padH, dilationW, dilationH, group, deformable_group,
+      im2col_step;
+  SSAttrs(attr)
+      .get<int>("kW", kW)
+      .get<int>("kH", kH)
+      .get<int>("dW", dW)
+      .get<int>("dH", dH)
+      .get<int>("padW", padW)
+      .get<int>("padH", padH)
+      .get<int>("dilationW", dilationW)
+      .get<int>("dilationH", dilationH)
+      .get<int>("group", group)
+      .get<int>("deformable_group", deformable_group)
+      .get<int>("im2col_step", im2col_step)
+      .done();
+
+  const auto& input = buildATensor(ctx, ins[0]);
+  const auto& weight = buildATensor(ctx, ins[1]);
+  const auto& offset = buildATensor(ctx, ins[2]);
+
+  auto output = buildATensor(ctx, outs[0]);
+  auto columns = buildATensor(ctx, outs[1]);
+  auto ones = buildATensor(ctx, outs[2]);
+
+  deform_conv_forward_cuda(input, weight, offset, output, columns, ones, kW, kH,
+                           dW, dH, padW, padH, dilationW, dilationH, group,
+                           deformable_group, im2col_step);
+}
+
+/*void deform_conv_backward_input_cuda(Tensor input, Tensor offset,
+ *                                     Tensor gradOutput, Tensor gradInput,
+ *                                     Tensor gradOffset, Tensor weight,
+ *                                     Tensor columns, int kW, int kH, int dW,
+ *                                     int dH, int padW, int padH, int
+ * dilationW, int dilationH, int group, int deformable_group, int im2col_step);
+ */
+void deform_conv_backward_input_cuda_parrots(CudaContext& ctx,
+                                             const SSElement& attr,
+                                             const OperatorBase::in_list_t& ins,
+                                             OperatorBase::out_list_t& outs) {
+  int kW, kH, dW, dH, padW, padH, dilationW, dilationH, group, deformable_group,
+      im2col_step;
+  SSAttrs(attr)
+      .get<int>("kW", kW)
+      .get<int>("kH", kH)
+      .get<int>("dW", dW)
+      .get<int>("dH", dH)
+      .get<int>("padW", padW)
+      .get<int>("padH", padH)
+      .get<int>("dilationW", dilationW)
+      .get<int>("dilationH", dilationH)
+      .get<int>("group", group)
+      .get<int>("deformable_group", deformable_group)
+      .get<int>("im2col_step", im2col_step)
+      .done();
+
+  const auto& input = buildATensor(ctx, ins[0]);
+  const auto& offset = buildATensor(ctx, ins[1]);
+  const auto& gradOutput = buildATensor(ctx, ins[2]);
+
+  auto gradInput = buildATensor(ctx, outs[0]);
+  auto gradOffset = buildATensor(ctx, outs[1]);
+  auto weight = buildATensor(ctx, outs[2]);
+  auto columns = buildATensor(ctx, outs[3]);
+
+  deform_conv_backward_input_cuda(input, offset, gradOutput, gradInput,
+                                  gradOffset, weight, columns, kW, kH, dW, dH,
+                                  padW, padH, dilationW, dilationH, group,
+                                  deformable_group, im2col_step);
+}
+
+/*void deform_conv_backward_parameters_cuda(
+ *     Tensor input, Tensor offset, Tensor gradOutput, Tensor gradWeight,
+ *     Tensor columns, Tensor ones, int kW, int kH, int dW, int dH, int padW,
+ *     int padH, int dilationW, int dilationH, int group, int deformable_group,
+ *     float scale, int im2col_step);
+ */
+void deform_conv_backward_parameters_cuda_parrots(
+    CudaContext& ctx, const SSElement& attr, const OperatorBase::in_list_t& ins,
+    OperatorBase::out_list_t& outs) {
+  int kW, kH, dW, dH, padW, padH, dilationW, dilationH, group, deformable_group,
+      im2col_step;
+  float scale;
+  SSAttrs(attr)
+      .get<int>("kW", kW)
+      .get<int>("kH", kH)
+      .get<int>("dW", dW)
+      .get<int>("dH", dH)
+      .get<int>("padW", padW)
+      .get<int>("padH", padH)
+      .get<int>("dilationW", dilationW)
+      .get<int>("dilationH", dilationH)
+      .get<int>("group", group)
+      .get<int>("deformable_group", deformable_group)
+      .get<float>("scale", scale)
+      .get<int>("im2col_step", im2col_step)
+      .done();
+
+  const auto& input = buildATensor(ctx, ins[0]);
+  const auto& offset = buildATensor(ctx, ins[1]);
+  const auto& gradOutput = buildATensor(ctx, ins[2]);
+
+  auto gradWeight = buildATensor(ctx, outs[0]);
+  auto columns = buildATensor(ctx, outs[1]);
+  auto ones = buildATensor(ctx, outs[2]);
+  deform_conv_backward_parameters_cuda(input, offset, gradOutput, gradWeight,
+                                       columns, ones, kW, kH, dW, dH, padW,
+                                       padH, dilationW, dilationH, group,
+                                       deformable_group, scale, im2col_step);
+}
+
+PARROTS_EXTENSION_REGISTER(deform_conv_forward)
+    .attr("kW")
+    .attr("kH")
+    .attr("dW")
+    .attr("dH")
+    .attr("padW")
+    .attr("padH")
+    .attr("dilationW")
+    .attr("dilationH")
+    .attr("group")
+    .attr("deformable_group")
+    .attr("im2col_step")
+    .input(3)
+    .output(3)
+    .apply(deform_conv_forward_cuda_parrots)
+    .done();
+
+PARROTS_EXTENSION_REGISTER(deform_conv_backward_input)
+    .attr("kW")
+    .attr("kH")
+    .attr("dW")
+    .attr("dH")
+    .attr("padW")
+    .attr("padH")
+    .attr("dilationW")
+    .attr("dilationH")
+    .attr("group")
+    .attr("deformable_group")
+    .attr("im2col_step")
+    .input(3)
+    .output(4)
+    .apply(deform_conv_backward_input_cuda_parrots)
+    .done();
+
+PARROTS_EXTENSION_REGISTER(deform_conv_backward_parameters)
+    .attr("kW")
+    .attr("kH")
+    .attr("dW")
+    .attr("dH")
+    .attr("padW")
+    .attr("padH")
+    .attr("dilationW")
+    .attr("dilationH")
+    .attr("group")
+    .attr("deformable_group")
+    .attr("scale")
+    .attr("im2col_step")
+    .input(3)
+    .output(3)
+    .apply(deform_conv_backward_parameters_cuda_parrots)
+    .done();

mmcv/ops/csrc/parrots/deform_conv_pytorch.h

+#ifndef DEFORM_CONV_PYTORCH_H
+#define DEFORM_CONV_PYTORCH_H
+#include <torch/extension.h>
+using namespace at;
+
+void deform_conv_forward_cuda(Tensor input, Tensor weight, Tensor offset,
+                              Tensor output, Tensor columns, Tensor ones,
+                              int kW, int kH, int dW, int dH, int padW,
+                              int padH, int dilationW, int dilationH, int group,
+                              int deformable_group, int im2col_step);
+
+void deform_conv_backward_input_cuda(Tensor input, Tensor offset,
+                                     Tensor gradOutput, Tensor gradInput,
+                                     Tensor gradOffset, Tensor weight,
+                                     Tensor columns, int kW, int kH, int dW,
+                                     int dH, int padW, int padH, int dilationW,
+                                     int dilationH, int group,
+                                     int deformable_group, int im2col_step);
+
+void deform_conv_backward_parameters_cuda(
+    Tensor input, Tensor offset, Tensor gradOutput, Tensor gradWeight,
+    Tensor columns, Tensor ones, int kW, int kH, int dW, int dH, int padW,
+    int padH, int dilationW, int dilationH, int group, int deformable_group,
+    float scale, int im2col_step);
+
+#endif  // DEFORM_CONV_PYTORCH_H

mmcv/ops/csrc/parrots/deform_roi_pool.cpp

-#include "parrots_cpp_helper.hpp"
+#include "pytorch_cpp_helper.hpp"
 
-void DeformRoIPoolForwardCUDAKernelLauncher(
-    const DArrayLite input, const DArrayLite rois, const DArrayLite offset,
-    DArrayLite output, int pooled_height, int pooled_width, float spatial_scale,
-    int sampling_ratio, float gamma, cudaStream_t stream);
+#ifdef MMCV_WITH_CUDA
+void DeformRoIPoolForwardCUDAKernelLauncher(Tensor input, Tensor rois,
+                                            Tensor offset, Tensor output,
+                                            int pooled_height, int pooled_width,
+                                            float spatial_scale,
+                                            int sampling_ratio, float gamma);
 
 void DeformRoIPoolBackwardCUDAKernelLauncher(
-    const DArrayLite grad_output, const DArrayLite input, const DArrayLite rois,
-    const DArrayLite offset, DArrayLite grad_input, DArrayLite grad_offset,
-    int pooled_height, int pooled_width, float spatial_scale,
-    int sampling_ratio, float gamma, cudaStream_t stream);
+    Tensor grad_output, Tensor input, Tensor rois, Tensor offset,
+    Tensor grad_input, Tensor grad_offset, int pooled_height, int pooled_width,
+    float spatial_scale, int sampling_ratio, float gamma);
 
-void deform_roi_pool_forward_cuda(CudaContext& ctx, const SSElement& attr,
-                                  const OperatorBase::in_list_t& ins,
-                                  OperatorBase::out_list_t& outs) {
-  int pooled_height;
-  int pooled_width;
-  float spatial_scale;
-  int sampling_ratio;
-  float gamma;
-  SSAttrs(attr)
-      .get<int>("pooled_height", pooled_height)
-      .get<int>("pooled_width", pooled_width)
-      .get<float>("spatial_scale", spatial_scale)
-      .get<int>("sampling_ratio", sampling_ratio)
-      .get<float>("gamma", gamma)
-      .done();
-
-  const auto& input = ins[0];
-  const auto& rois = ins[1];
-  const auto& offset = ins[2];
-
-  auto& output = outs[0];
-
-  cudaStream_t stream = getStreamNative<CudaDevice>(ctx.getStream());
-  DeformRoIPoolForwardCUDAKernelLauncher(
-      input, rois, offset, output, pooled_height, pooled_width, spatial_scale,
-      sampling_ratio, gamma, stream);
+void deform_roi_pool_forward_cuda(Tensor input, Tensor rois, Tensor offset,
+                                  Tensor output, int pooled_height,
+                                  int pooled_width, float spatial_scale,
+                                  int sampling_ratio, float gamma) {
+  DeformRoIPoolForwardCUDAKernelLauncher(input, rois, offset, output,
+                                         pooled_height, pooled_width,
+                                         spatial_scale, sampling_ratio, gamma);
 }
 
-void deform_roi_pool_backward_cuda(CudaContext& ctx, const SSElement& attr,
-                                   const OperatorBase::in_list_t& ins,
-                                   OperatorBase::out_list_t& outs) {
-  int pooled_height;
-  int pooled_width;
-  float spatial_scale;
-  int sampling_ratio;
-  float gamma;
-
-  SSAttrs(attr)
-      .get<int>("pooled_height", pooled_height)
-      .get<int>("pooled_width", pooled_width)
-      .get<float>("spatial_scale", spatial_scale)
-      .get<int>("sampling_ratio", sampling_ratio)
-      .get<float>("gamma", gamma)
-      .done();
-
-  const auto& grad_output = ins[0];
-  const auto& input = ins[1];
-  const auto& rois = ins[2];
-  const auto& offset = ins[3];
-
-  auto& grad_input = outs[0];
-  auto& grad_offset = outs[1];
-
-  cudaStream_t stream = getStreamNative<CudaDevice>(ctx.getStream());
+void deform_roi_pool_backward_cuda(Tensor grad_output, Tensor input,
+                                   Tensor rois, Tensor offset,
+                                   Tensor grad_input, Tensor grad_offset,
+                                   int pooled_height, int pooled_width,
+                                   float spatial_scale, int sampling_ratio,
+                                   float gamma) {
   DeformRoIPoolBackwardCUDAKernelLauncher(
       grad_output, input, rois, offset, grad_input, grad_offset, pooled_height,
-      pooled_width, spatial_scale, sampling_ratio, gamma, stream);
+      pooled_width, spatial_scale, sampling_ratio, gamma);
 }
+#endif
+
+void deform_roi_pool_forward(Tensor input, Tensor rois, Tensor offset,
+                             Tensor output, int pooled_height, int pooled_width,
+                             float spatial_scale, int sampling_ratio,
+                             float gamma) {
+  if (input.device().is_cuda()) {
+#ifdef MMCV_WITH_CUDA
+    CHECK_CUDA_INPUT(input);
+    CHECK_CUDA_INPUT(rois);
+    CHECK_CUDA_INPUT(offset);
+    CHECK_CUDA_INPUT(output);
 
-PARROTS_EXTENSION_REGISTER(deform_roi_pool_forward)
-    .attr("pooled_height")
-    .attr("pooled_width")
-    .attr("spatial_scale")
-    .attr("sampling_ratio")
-    .attr("gamma")
-    .input(3)
-    .output(1)
-    .apply(deform_roi_pool_forward_cuda)
-    .done();
+    deform_roi_pool_forward_cuda(input, rois, offset, output, pooled_height,
+                                 pooled_width, spatial_scale, sampling_ratio,
+                                 gamma);
+#else
+    AT_ERROR("DeformRoIPool is not compiled with GPU support");
+#endif
+  } else {
+    AT_ERROR("DeformRoIPool is not implemented on CPU");
+  }
+}
 
-PARROTS_EXTENSION_REGISTER(deform_roi_pool_backward)
-    .attr("pooled_height")
-    .attr("pooled_width")
-    .attr("spatial_scale")
-    .attr("sampling_ratio")
-    .attr("gamma")
-    .input(4)
-    .output(2)
-    .apply(deform_roi_pool_backward_cuda)
-    .done();
+void deform_roi_pool_backward(Tensor grad_output, Tensor input, Tensor rois,
+                              Tensor offset, Tensor grad_input,
+                              Tensor grad_offset, int pooled_height,
+                              int pooled_width, float spatial_scale,
+                              int sampling_ratio, float gamma) {
+  if (grad_output.device().is_cuda()) {
+#ifdef MMCV_WITH_CUDA
+    CHECK_CUDA_INPUT(grad_output);
+    CHECK_CUDA_INPUT(input);
+    CHECK_CUDA_INPUT(rois);
+    CHECK_CUDA_INPUT(offset);
+    CHECK_CUDA_INPUT(grad_input);
+    CHECK_CUDA_INPUT(grad_offset);
+
+    deform_roi_pool_backward_cuda(grad_output, input, rois, offset, grad_input,
+                                  grad_offset, pooled_height, pooled_width,
+                                  spatial_scale, sampling_ratio, gamma);
+#else
+    AT_ERROR("DeformRoIPool is not compiled with GPU support");
+#endif
+  } else {
+    AT_ERROR("DeformRoIPool is not implemented on CPU");
+  }
+}

mmcv/ops/csrc/parrots/deform_roi_pool_cuda.cu

 #include "deform_roi_pool_cuda_kernel.cuh"
-#include "parrots_cuda_helper.hpp"
+#include "pytorch_cuda_helper.hpp"
 
-void DeformRoIPoolForwardCUDAKernelLauncher(
-    const DArrayLite input, const DArrayLite rois, const DArrayLite offset,
-    DArrayLite output, int pooled_height, int pooled_width, float spatial_scale,
-    int sampling_ratio, float gamma, cudaStream_t stream) {
-  int output_size = output.size();
-  int channels = input.dim(1);
-  int height = input.dim(2);
-  int width = input.dim(3);
+void DeformRoIPoolForwardCUDAKernelLauncher(Tensor input, Tensor rois,
+                                            Tensor offset, Tensor output,
+                                            int pooled_height, int pooled_width,
+                                            float spatial_scale,
+                                            int sampling_ratio, float gamma) {
+  int output_size = output.numel();
+  int channels = input.size(1);
+  int height = input.size(2);
+  int width = input.size(3);
 
-  PARROTS_DISPATCH_FLOATING_TYPES_AND_HALF(
-      input.elemType().prim(), ([&] {
+  at::cuda::CUDAGuard device_guard(input.device());
+  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+      input.scalar_type(), "deform_roi_pool_forward_cuda_kernel", [&] {
         deform_roi_pool_forward_cuda_kernel<scalar_t>
             <<<GET_BLOCKS(output_size), THREADS_PER_BLOCK, 0, stream>>>(
-                output_size, input.ptr<scalar_t>(), rois.ptr<scalar_t>(),
-                offset.ptr<scalar_t>(), output.ptr<scalar_t>(), pooled_height,
-                pooled_width, spatial_scale, sampling_ratio, gamma, channels,
-                height, width);
-      }));
+                output_size, input.data_ptr<scalar_t>(),
+                rois.data_ptr<scalar_t>(), offset.data_ptr<scalar_t>(),
+                output.data_ptr<scalar_t>(), pooled_height, pooled_width,
+                static_cast<scalar_t>(spatial_scale), sampling_ratio,
+                static_cast<scalar_t>(gamma), channels, height, width);
+      });
 
-  PARROTS_CUDA_CHECK(cudaGetLastError());
+  AT_CUDA_CHECK(cudaGetLastError());
 }
 
 void DeformRoIPoolBackwardCUDAKernelLauncher(
-    const DArrayLite grad_output, const DArrayLite input, const DArrayLite rois,
-    const DArrayLite offset, DArrayLite grad_input, DArrayLite grad_offset,
-    int pooled_height, int pooled_width, float spatial_scale,
-    int sampling_ratio, float gamma, cudaStream_t stream) {
-  int output_size = grad_output.size();
-  int channels = grad_input.dim(1);
-  int height = grad_input.dim(2);
-  int width = grad_input.dim(3);
+    Tensor grad_output, Tensor input, Tensor rois, Tensor offset,
+    Tensor grad_input, Tensor grad_offset, int pooled_height, int pooled_width,
+    float spatial_scale, int sampling_ratio, float gamma) {
+  int output_size = grad_output.numel();
+  int channels = grad_input.size(1);
+  int height = grad_input.size(2);
+  int width = grad_input.size(3);
 
-  PARROTS_DISPATCH_FLOATING_TYPES_AND_HALF(
-      grad_output.elemType().prim(), ([&] {
+  at::cuda::CUDAGuard device_guard(grad_output.device());
+  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+      grad_output.scalar_type(), "deform_roi_pool_backward_cuda_kernel", [&] {
         deform_roi_pool_backward_cuda_kernel<scalar_t>
             <<<GET_BLOCKS(output_size), THREADS_PER_BLOCK, 0, stream>>>(
-                output_size, grad_output.ptr<scalar_t>(), input.ptr<scalar_t>(),
-                rois.ptr<scalar_t>(), offset.ptr<scalar_t>(),
-                grad_input.ptr<scalar_t>(), grad_offset.ptr<scalar_t>(),
-                pooled_height, pooled_width, spatial_scale, sampling_ratio,
-                gamma, channels, height, width);
-      }));
+                output_size, grad_output.data_ptr<scalar_t>(),
+                input.data_ptr<scalar_t>(), rois.data_ptr<scalar_t>(),
+                offset.data_ptr<scalar_t>(), grad_input.data_ptr<scalar_t>(),
+                grad_offset.data_ptr<scalar_t>(), pooled_height, pooled_width,
+                static_cast<scalar_t>(spatial_scale), sampling_ratio,
+                static_cast<scalar_t>(gamma), channels, height, width);
+      });
 
-  PARROTS_CUDA_CHECK(cudaGetLastError());
+  AT_CUDA_CHECK(cudaGetLastError());
 }

mmcv/ops/csrc/parrots/deform_roi_pool_parrots.cpp

+#include <parrots/compute/aten.hpp>
+#include <parrots/extension.hpp>
+#include <parrots/foundation/ssattrs.hpp>
+
+#include "deform_roi_pool_pytorch.h"
+
+using namespace parrots;
+
+/*void deform_roi_pool_forward_cuda(Tensor input, Tensor rois, Tensor offset,
+ *                                  Tensor output, int pooled_height,
+ *                                  int pooled_width, float spatial_scale,
+ *                                  int sampling_ratio, float gamma);
+ */
+void deform_roi_pool_forward_cuda_parrots(CudaContext& ctx,
+                                          const SSElement& attr,
+                                          const OperatorBase::in_list_t& ins,
+                                          OperatorBase::out_list_t& outs) {
+  int pooled_height;
+  int pooled_width;
+  float spatial_scale;
+  int sampling_ratio;
+  float gamma;
+  SSAttrs(attr)
+      .get<int>("pooled_height", pooled_height)
+      .get<int>("pooled_width", pooled_width)
+      .get<float>("spatial_scale", spatial_scale)
+      .get<int>("sampling_ratio", sampling_ratio)
+      .get<float>("gamma", gamma)
+      .done();
+
+  const auto& input = buildATensor(ctx, ins[0]);
+  const auto& rois = buildATensor(ctx, ins[1]);
+  const auto& offset = buildATensor(ctx, ins[2]);
+
+  auto output = buildATensor(ctx, outs[0]);
+  deform_roi_pool_forward_cuda(input, rois, offset, output, pooled_height,
+                               pooled_width, spatial_scale, sampling_ratio,
+                               gamma);
+}
+
+/*void deform_roi_pool_backward_cuda(Tensor grad_output, Tensor input,
+ *                                   Tensor rois, Tensor offset,
+ *                                   Tensor grad_input, Tensor grad_offset,
+ *                                   int pooled_height, int pooled_width,
+ *                                   float spatial_scale, int sampling_ratio,
+ *                                   float gamma);
+ */
+void deform_roi_pool_backward_cuda_parrots(CudaContext& ctx,
+                                           const SSElement& attr,
+                                           const OperatorBase::in_list_t& ins,
+                                           OperatorBase::out_list_t& outs) {
+  int pooled_height;
+  int pooled_width;
+  float spatial_scale;
+  int sampling_ratio;
+  float gamma;
+
+  SSAttrs(attr)
+      .get<int>("pooled_height", pooled_height)
+      .get<int>("pooled_width", pooled_width)
+      .get<float>("spatial_scale", spatial_scale)
+      .get<int>("sampling_ratio", sampling_ratio)
+      .get<float>("gamma", gamma)
+      .done();
+
+  const auto& grad_output = buildATensor(ctx, ins[0]);
+  const auto& input = buildATensor(ctx, ins[1]);
+  const auto& rois = buildATensor(ctx, ins[2]);
+  const auto& offset = buildATensor(ctx, ins[3]);
+
+  auto grad_input = buildATensor(ctx, outs[0]);
+  auto grad_offset = buildATensor(ctx, outs[1]);
+
+  deform_roi_pool_backward_cuda(grad_output, input, rois, offset, grad_input,
+                                grad_offset, pooled_height, pooled_width,
+                                spatial_scale, sampling_ratio, gamma);
+}
+
+PARROTS_EXTENSION_REGISTER(deform_roi_pool_forward)
+    .attr("pooled_height")
+    .attr("pooled_width")
+    .attr("spatial_scale")
+    .attr("sampling_ratio")
+    .attr("gamma")
+    .input(3)
+    .output(1)
+    .apply(deform_roi_pool_forward_cuda_parrots)
+    .done();
+
+PARROTS_EXTENSION_REGISTER(deform_roi_pool_backward)
+    .attr("pooled_height")
+    .attr("pooled_width")
+    .attr("spatial_scale")
+    .attr("sampling_ratio")
+    .attr("gamma")
+    .input(4)
+    .output(2)
+    .apply(deform_roi_pool_backward_cuda_parrots)
+    .done();

mmcv/ops/csrc/parrots/deform_roi_pool_pytorch.h

+#ifndef DEFORM_ROI_POOL_PYTORCH_H
+#define DEFORM_ROI_POOL_PYTORCH_H
+#include <torch/extension.h>
+using namespace at;
+
+void deform_roi_pool_forward_cuda(Tensor input, Tensor rois, Tensor offset,
+                                  Tensor output, int pooled_height,
+                                  int pooled_width, float spatial_scale,
+                                  int sampling_ratio, float gamma);
+
+void deform_roi_pool_backward_cuda(Tensor grad_output, Tensor input,
+                                   Tensor rois, Tensor offset,
+                                   Tensor grad_input, Tensor grad_offset,
+                                   int pooled_height, int pooled_width,
+                                   float spatial_scale, int sampling_ratio,
+                                   float gamma);
+#endif  // DEFORM_ROI_POOL_PYTORCH_H

mmcv/ops/csrc/parrots/focal_loss.cpp

-// Copyright (c) 2018, SenseTime.
-#include "parrots_cpp_helper.hpp"
-
-void SigmoidFocalLossForwardCUDAKernelLauncher(
-    const DArrayLite input, const DArrayLite target, const DArrayLite weight,
-    DArrayLite output, float gamma, float alpha, cudaStream_t stream);
-
-void SigmoidFocalLossBackwardCUDAKernelLauncher(
-    const DArrayLite input, const DArrayLite target, const DArrayLite weight,
-    DArrayLite grad_input, float gamma, float alpha, cudaStream_t stream);
-
-void SoftmaxFocalLossForwardCUDAKernelLauncher(
-    const DArrayLite input, const DArrayLite target, const DArrayLite weight,
-    DArrayLite output, float gamma, float alpha, cudaStream_t stream);
-
-void SoftmaxFocalLossBackwardCUDAKernelLauncher(
-    const DArrayLite input, const DArrayLite target, const DArrayLite weight,
-    DArrayLite buff, DArrayLite grad_input, float gamma, float alpha,
-    cudaStream_t stream);
-
-void sigmoid_focal_loss_forward_cuda(CudaContext& ctx, const SSElement& attr,
-                                     const OperatorBase::in_list_t& ins,
-                                     OperatorBase::out_list_t& outs) {
-  float gamma;
-  float alpha;
-  SSAttrs(attr).get<float>("gamma", gamma).get<float>("alpha", alpha).done();
-
-  // get inputs and outputs
-  const auto& input = ins[0];
-  const auto& target = ins[1];
-  const auto& weight = ins[2];
-
-  auto& output = outs[0];
-
-  cudaStream_t stream = getStreamNative<CudaDevice>(ctx.getStream());
-
+#include "pytorch_cpp_helper.hpp"
+
+#ifdef MMCV_WITH_CUDA
+void SigmoidFocalLossForwardCUDAKernelLauncher(Tensor input, Tensor target,
+                                               Tensor weight, Tensor output,
+                                               const float gamma,
+                                               const float alpha);
+
+void SigmoidFocalLossBackwardCUDAKernelLauncher(Tensor input, Tensor target,
+                                                Tensor weight,
+                                                Tensor grad_input,
+                                                const float gamma,
+                                                const float alpha);
+
+void SoftmaxFocalLossForwardCUDAKernelLauncher(Tensor input, Tensor target,
+                                               Tensor weight, Tensor output,
+                                               const float gamma,
+                                               const float alpha);
+
+void SoftmaxFocalLossBackwardCUDAKernelLauncher(Tensor input, Tensor target,
+                                                Tensor weight, Tensor buff,
+                                                Tensor grad_input,
+                                                const float gamma,
+                                                const float alpha);
+
+void sigmoid_focal_loss_forward_cuda(Tensor input, Tensor target, Tensor weight,
+                                     Tensor output, float gamma, float alpha) {
   SigmoidFocalLossForwardCUDAKernelLauncher(input, target, weight, output,
-                                            gamma, alpha, stream);
+                                            gamma, alpha);
 }
 
-void sigmoid_focal_loss_backward_cuda(CudaContext& ctx, const SSElement& attr,
-                                      const OperatorBase::in_list_t& ins,
-                                      OperatorBase::out_list_t& outs) {
-  float gamma;
-  float alpha;
-  SSAttrs(attr).get<float>("gamma", gamma).get<float>("alpha", alpha).done();
-
-  // get inputs and outputs
-  const auto& input = ins[0];
-  const auto& target = ins[1];
-  const auto& weight = ins[2];
-
-  auto& grad_input = outs[0];
-
-  cudaStream_t stream = getStreamNative<CudaDevice>(ctx.getStream());
+void sigmoid_focal_loss_backward_cuda(Tensor input, Tensor target,
+                                      Tensor weight, Tensor grad_input,
+                                      float gamma, float alpha) {
   SigmoidFocalLossBackwardCUDAKernelLauncher(input, target, weight, grad_input,
-                                             gamma, alpha, stream);
+                                             gamma, alpha);
 }
 
-void softmax_focal_loss_forward_cuda(CudaContext& ctx, const SSElement& attr,
-                                     const OperatorBase::in_list_t& ins,
-                                     OperatorBase::out_list_t& outs) {
-  float gamma;
-  float alpha;
-  SSAttrs(attr).get<float>("gamma", gamma).get<float>("alpha", alpha).done();
-
-  // get inputs and outputs
-  const auto& input = ins[0];
-  const auto& target = ins[1];
-  const auto& weight = ins[2];
-
-  auto& grad_input = outs[0];
-
-  cudaStream_t stream = getStreamNative<CudaDevice>(ctx.getStream());
-
-  SoftmaxFocalLossForwardCUDAKernelLauncher(input, target, weight, grad_input,
-                                            gamma, alpha, stream);
+void softmax_focal_loss_forward_cuda(Tensor input, Tensor target, Tensor weight,
+                                     Tensor output, float gamma, float alpha) {
+  SoftmaxFocalLossForwardCUDAKernelLauncher(input, target, weight, output,
+                                            gamma, alpha);
 }
 
-void softmax_focal_loss_backward_cuda(CudaContext& ctx, const SSElement& attr,
-                                      const OperatorBase::in_list_t& ins,
-                                      OperatorBase::out_list_t& outs) {
-  float gamma;
-  float alpha;
-  SSAttrs(attr).get<float>("gamma", gamma).get<float>("alpha", alpha).done();
-
-  // get inputs and outputs
-  const auto& input = ins[0];
-  const auto& target = ins[1];
-  const auto& weight = ins[2];
+void softmax_focal_loss_backward_cuda(Tensor input, Tensor target,
+                                      Tensor weight, Tensor buff,
+                                      Tensor grad_input, float gamma,
+                                      float alpha) {
+  SoftmaxFocalLossBackwardCUDAKernelLauncher(input, target, weight, buff,
+                                             grad_input, gamma, alpha);
+}
+#endif
+
+void sigmoid_focal_loss_forward(Tensor input, Tensor target, Tensor weight,
+                                Tensor output, float gamma, float alpha) {
+  if (input.device().is_cuda()) {
+#ifdef MMCV_WITH_CUDA
+    CHECK_CUDA_INPUT(input);
+    CHECK_CUDA_INPUT(target);
+    CHECK_CUDA_INPUT(weight);
+    CHECK_CUDA_INPUT(output);
+
+    sigmoid_focal_loss_forward_cuda(input, target, weight, output, gamma,
+                                    alpha);
+#else
+    AT_ERROR("SigmoidFocalLoss is not compiled with GPU support");
+#endif
+  } else {
+    AT_ERROR("SigmoidFocalLoss is not implemented on CPU");
+  }
+}
 
-  auto& buff = outs[0];
-  auto& grad_input = outs[1];
+void sigmoid_focal_loss_backward(Tensor input, Tensor target, Tensor weight,
+                                 Tensor grad_input, float gamma, float alpha) {
+  if (input.device().is_cuda()) {
+#ifdef MMCV_WITH_CUDA
+    CHECK_CUDA_INPUT(input);
+    CHECK_CUDA_INPUT(target);
+    CHECK_CUDA_INPUT(weight);
+    CHECK_CUDA_INPUT(grad_input);
+
+    sigmoid_focal_loss_backward_cuda(input, target, weight, grad_input, gamma,
+                                     alpha);
+#else
+    AT_ERROR("SigmoidFocalLoss is not compiled with GPU support");
+#endif
+  } else {
+    AT_ERROR("SigmoidFocalLoss is not implemented on CPU");
+  }
+}
 
-  cudaStream_t stream = getStreamNative<CudaDevice>(ctx.getStream());
-  SoftmaxFocalLossBackwardCUDAKernelLauncher(input, target, weight, buff,
-                                             grad_input, gamma, alpha, stream);
+void softmax_focal_loss_forward(Tensor input, Tensor target, Tensor weight,
+                                Tensor output, float gamma, float alpha) {
+  if (input.device().is_cuda()) {
+#ifdef MMCV_WITH_CUDA
+    CHECK_CUDA_INPUT(input);
+    CHECK_CUDA_INPUT(target);
+    CHECK_CUDA_INPUT(weight);
+    CHECK_CUDA_INPUT(output);
+
+    softmax_focal_loss_forward_cuda(input, target, weight, output, gamma,
+                                    alpha);
+#else
+    AT_ERROR("SoftmaxFocalLoss is not compiled with GPU support");
+#endif
+  } else {
+    AT_ERROR("SoftmaxFocalLoss is not implemented on CPU");
+  }
 }
 
-PARROTS_EXTENSION_REGISTER(sigmoid_focal_loss_forward)
-    .attr("gamma")
-    .attr("alpha")
-    .input(3)
-    .output(1)
-    .apply(sigmoid_focal_loss_forward_cuda)
-    .done();
-
-PARROTS_EXTENSION_REGISTER(sigmoid_focal_loss_backward)
-    .attr("gamma")
-    .attr("alpha")
-    .input(3)
-    .output(1)
-    .apply(sigmoid_focal_loss_backward_cuda)
-    .done();
-
-PARROTS_EXTENSION_REGISTER(softmax_focal_loss_forward)
-    .attr("gamma")
-    .attr("alpha")
-    .input(3)
-    .output(1)
-    .apply(softmax_focal_loss_forward_cuda)
-    .done();
-
-PARROTS_EXTENSION_REGISTER(softmax_focal_loss_backward)
-    .attr("gamma")
-    .attr("alpha")
-    .input(3)
-    .output(2)
-    .apply(softmax_focal_loss_backward_cuda)
-    .done();
+void softmax_focal_loss_backward(Tensor input, Tensor target, Tensor weight,
+                                 Tensor buff, Tensor grad_input, float gamma,
+                                 float alpha) {
+  if (input.device().is_cuda()) {
+#ifdef MMCV_WITH_CUDA
+    CHECK_CUDA_INPUT(input);
+    CHECK_CUDA_INPUT(target);
+    CHECK_CUDA_INPUT(weight);
+    CHECK_CUDA_INPUT(buff);
+    CHECK_CUDA_INPUT(grad_input);
+
+    softmax_focal_loss_backward_cuda(input, target, weight, buff, grad_input,
+                                     gamma, alpha);
+#else
+    AT_ERROR("SoftmaxFocalLoss is not compiled with GPU support");
+#endif
+  } else {
+    AT_ERROR("SoftmaxFocalLoss is not implemented on CPU");
+  }
+}

mmcv/ops/csrc/parrots/focal_loss_cuda.cu

-#include "parrots_cuda_helper.hpp"
+#include "pytorch_cuda_helper.hpp"
 #include "sigmoid_focal_loss_cuda_kernel.cuh"
 #include "softmax_focal_loss_cuda_kernel.cuh"
 
-void SigmoidFocalLossForwardCUDAKernelLauncher(
-    const DArrayLite input, const DArrayLite target, const DArrayLite weight,
-    DArrayLite output, float gamma, float alpha, cudaStream_t stream) {
-  int output_size = output.size();
-  int num_classes = input.dim(1);
-
-  PARROTS_DISPATCH_FLOATING_TYPES_AND_HALF(
-      input.elemType().prim(), ([&] {
+void SigmoidFocalLossForwardCUDAKernelLauncher(Tensor input, Tensor target,
+                                               Tensor weight, Tensor output,
+                                               const float gamma,
+                                               const float alpha) {
+  int output_size = output.numel();
+  int num_classes = input.size(1);
+  AT_ASSERTM(target.max().item<int64_t>() <= (int64_t)num_classes,
+             "target label should smaller or equal than num classes");
+  at::cuda::CUDAGuard device_guard(input.device());
+  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+      input.scalar_type(), "sigmoid_focal_loss_forward_cuda_kernel", [&] {
         sigmoid_focal_loss_forward_cuda_kernel<scalar_t>
             <<<GET_BLOCKS(output_size), THREADS_PER_BLOCK, 0, stream>>>(
-                output_size, input.ptr<scalar_t>(), target.ptr<int64_t>(),
-                weight.ptr<scalar_t>(), output.ptr<scalar_t>(), gamma, alpha,
-                num_classes);
-      }));
+                output_size, input.data_ptr<scalar_t>(),
+                target.data_ptr<int64_t>(), weight.data_ptr<scalar_t>(),
+                output.data_ptr<scalar_t>(), gamma, alpha, num_classes);
+      });
 
-  PARROTS_CUDA_CHECK(cudaGetLastError());
+  AT_CUDA_CHECK(cudaGetLastError());
 }
 
-void SigmoidFocalLossBackwardCUDAKernelLauncher(
-    const DArrayLite input, const DArrayLite target, const DArrayLite weight,
-    DArrayLite grad_input, float gamma, float alpha, cudaStream_t stream) {
-  int output_size = grad_input.size();
-  int num_classes = input.dim(1);
+void SigmoidFocalLossBackwardCUDAKernelLauncher(Tensor input, Tensor target,
+                                                Tensor weight,
+                                                Tensor grad_input,
+                                                const float gamma,
+                                                const float alpha) {
+  int output_size = grad_input.numel();
+  int num_classes = input.size(1);
 
-  PARROTS_DISPATCH_FLOATING_TYPES_AND_HALF(
-      input.elemType().prim(), ([&] {
+  at::cuda::CUDAGuard device_guard(grad_input.device());
+  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+      input.scalar_type(), "sigmoid_focal_loss_backward_cuda_kernel", [&] {
         sigmoid_focal_loss_backward_cuda_kernel<scalar_t>
             <<<GET_BLOCKS(output_size), THREADS_PER_BLOCK, 0, stream>>>(
-                output_size, input.ptr<scalar_t>(), target.ptr<int64_t>(),
-                weight.ptr<scalar_t>(), grad_input.ptr<scalar_t>(), gamma,
-                alpha, num_classes);
-      }));
+                output_size, input.data_ptr<scalar_t>(),
+                target.data_ptr<int64_t>(), weight.data_ptr<scalar_t>(),
+                grad_input.data_ptr<scalar_t>(), gamma, alpha, num_classes);
+      });
 
-  PARROTS_CUDA_CHECK(cudaGetLastError());
+  AT_CUDA_CHECK(cudaGetLastError());
 }
 
-void SoftmaxFocalLossForwardCUDAKernelLauncher(
-    const DArrayLite softmax, const DArrayLite target, const DArrayLite weight,
-    DArrayLite output, float gamma, float alpha, cudaStream_t stream) {
-  int output_size = output.size();
-  int num_classes = softmax.dim(1);
+void SoftmaxFocalLossForwardCUDAKernelLauncher(Tensor softmax, Tensor target,
+                                               Tensor weight, Tensor output,
+                                               const float gamma,
+                                               const float alpha) {
+  int output_size = output.numel();
+  int num_classes = softmax.size(1);
 
-  PARROTS_DISPATCH_FLOATING_TYPES_AND_HALF(
-      softmax.elemType().prim(), ([&] {
+  AT_ASSERTM(target.max().item<int64_t>() <= (int64_t)num_classes,
+             "target label should smaller or equal than num classes");
+  at::cuda::CUDAGuard device_guard(softmax.device());
+  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+      softmax.scalar_type(), "softmax_focal_loss_forward_cuda_kernel", [&] {
         softmax_focal_loss_forward_cuda_kernel<scalar_t>
             <<<GET_BLOCKS(output_size), THREADS_PER_BLOCK, 0, stream>>>(
-                output_size, softmax.ptr<scalar_t>(), target.ptr<int64_t>(),
-                weight.ptr<scalar_t>(), output.ptr<scalar_t>(), gamma, alpha,
-                num_classes);
-      }));
+                output_size, softmax.data_ptr<scalar_t>(),
+                target.data_ptr<int64_t>(), weight.data_ptr<scalar_t>(),
+                output.data_ptr<scalar_t>(), gamma, alpha, num_classes);
+      });
 
-  PARROTS_CUDA_CHECK(cudaGetLastError());
+  AT_CUDA_CHECK(cudaGetLastError());
 }
 
-void SoftmaxFocalLossBackwardCUDAKernelLauncher(
-    const DArrayLite softmax, const DArrayLite target, const DArrayLite weight,
-    DArrayLite buff, DArrayLite grad_input, float gamma, float alpha,
-    cudaStream_t stream) {
-  int output_size = buff.size();
-  int num_classes = softmax.dim(1);
+void SoftmaxFocalLossBackwardCUDAKernelLauncher(Tensor softmax, Tensor target,
+                                                Tensor weight, Tensor buff,
+                                                Tensor grad_input,
+                                                const float gamma,
+                                                const float alpha) {
+  int num_classes = softmax.size(1);
 
-  PARROTS_DISPATCH_FLOATING_TYPES_AND_HALF(
-      grad_input.elemType().prim(), ([&] {
+  int output_size = buff.numel();
+  at::cuda::CUDAGuard device_guard(grad_input.device());
+  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+      grad_input.scalar_type(),
+      "softmax_focal_loss_backward_cuda1_"
+      "kernel",
+      [&] {
         softmax_focal_loss_backward_cuda1_kernel<scalar_t>
             <<<GET_BLOCKS(output_size), THREADS_PER_BLOCK, 0, stream>>>(
-                output_size, softmax.ptr<scalar_t>(), target.ptr<int64_t>(),
-                weight.ptr<scalar_t>(), buff.ptr<scalar_t>(), gamma, alpha,
-                num_classes);
-      }));
-
-  PARROTS_CUDA_CHECK(cudaGetLastError());
+                output_size, softmax.data_ptr<scalar_t>(),
+                target.data_ptr<int64_t>(), weight.data_ptr<scalar_t>(),
+                buff.data_ptr<scalar_t>(), gamma, alpha, num_classes);
+      });
 
-  output_size = grad_input.size();
+  AT_CUDA_CHECK(cudaGetLastError());
 
-  PARROTS_DISPATCH_FLOATING_TYPES_AND_HALF(
-      grad_input.elemType().prim(), ([&] {
+  output_size = grad_input.numel();
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+      grad_input.scalar_type(),
+      "softmax_focal_loss_backward_cuda2_"
+      "kernel",
+      [&] {
         softmax_focal_loss_backward_cuda2_kernel<scalar_t>
             <<<GET_BLOCKS(output_size), THREADS_PER_BLOCK, 0, stream>>>(
-                output_size, softmax.ptr<scalar_t>(), target.ptr<int64_t>(),
-                buff.ptr<scalar_t>(), grad_input.ptr<scalar_t>(), num_classes);
-      }));
+                output_size, softmax.data_ptr<scalar_t>(),
+                target.data_ptr<int64_t>(), buff.data_ptr<scalar_t>(),
+                grad_input.data_ptr<scalar_t>(), num_classes);
+      });
 
-  PARROTS_CUDA_CHECK(cudaGetLastError());
+  AT_CUDA_CHECK(cudaGetLastError());
 }

mmcv/ops/csrc/parrots/focal_loss_parrots.cpp

+#include <parrots/compute/aten.hpp>
+#include <parrots/extension.hpp>
+#include <parrots/foundation/ssattrs.hpp>
+
+#include "focal_loss_pytorch.h"
+
+using namespace parrots;
+
+void sigmoid_focal_loss_forward_cuda_parrots(CudaContext& ctx,
+                                             const SSElement& attr,
+                                             const OperatorBase::in_list_t& ins,
+                                             OperatorBase::out_list_t& outs) {
+  float gamma;
+  float alpha;
+  SSAttrs(attr).get<float>("gamma", gamma).get<float>("alpha", alpha).done();
+
+  // get inputs and outputs
+  const auto& input = buildATensor(ctx, ins[0]);
+  const auto& target = buildATensor(ctx, ins[1]);
+  const auto& weight = buildATensor(ctx, ins[2]);
+
+  auto output = buildATensor(ctx, outs[0]);
+
+  sigmoid_focal_loss_forward_cuda(input, target, weight, output, gamma, alpha);
+}
+
+void sigmoid_focal_loss_backward_cuda_parrots(
+    CudaContext& ctx, const SSElement& attr, const OperatorBase::in_list_t& ins,
+    OperatorBase::out_list_t& outs) {
+  float gamma;
+  float alpha;
+  SSAttrs(attr).get<float>("gamma", gamma).get<float>("alpha", alpha).done();
+
+  // get inputs and outputs
+  const auto& input = buildATensor(ctx, ins[0]);
+  const auto& target = buildATensor(ctx, ins[1]);
+  const auto& weight = buildATensor(ctx, ins[2]);
+
+  auto grad_input = buildATensor(ctx, outs[0]);
+
+  sigmoid_focal_loss_backward_cuda(input, target, weight, grad_input, gamma,
+                                   alpha);
+}
+
+void softmax_focal_loss_forward_cuda_parrots(CudaContext& ctx,
+                                             const SSElement& attr,
+                                             const OperatorBase::in_list_t& ins,
+                                             OperatorBase::out_list_t& outs) {
+  float gamma;
+  float alpha;
+  SSAttrs(attr).get<float>("gamma", gamma).get<float>("alpha", alpha).done();
+
+  // get inputs and outputs
+  const auto& input = buildATensor(ctx, ins[0]);
+  const auto& target = buildATensor(ctx, ins[1]);
+  const auto& weight = buildATensor(ctx, ins[2]);
+
+  auto output = buildATensor(ctx, outs[0]);
+  softmax_focal_loss_forward_cuda(input, target, weight, output, gamma, alpha);
+}
+
+void softmax_focal_loss_backward_cuda_parrots(
+    CudaContext& ctx, const SSElement& attr, const OperatorBase::in_list_t& ins,
+    OperatorBase::out_list_t& outs) {
+  float gamma;
+  float alpha;
+  SSAttrs(attr).get<float>("gamma", gamma).get<float>("alpha", alpha).done();
+
+  // get inputs and outputs
+  const auto& input = buildATensor(ctx, ins[0]);
+  const auto& target = buildATensor(ctx, ins[1]);
+  const auto& weight = buildATensor(ctx, ins[2]);
+
+  auto buff = buildATensor(ctx, outs[0]);
+  auto grad_input = buildATensor(ctx, outs[1]);
+  softmax_focal_loss_backward_cuda(input, target, weight, buff, grad_input,
+                                   gamma, alpha);
+}
+
+PARROTS_EXTENSION_REGISTER(sigmoid_focal_loss_forward)
+    .attr("gamma")
+    .attr("alpha")
+    .input(3)
+    .output(1)
+    .apply(sigmoid_focal_loss_forward_cuda_parrots)
+    .done();
+
+PARROTS_EXTENSION_REGISTER(sigmoid_focal_loss_backward)
+    .attr("gamma")
+    .attr("alpha")
+    .input(3)
+    .output(1)
+    .apply(sigmoid_focal_loss_backward_cuda_parrots)
+    .done();
+
+PARROTS_EXTENSION_REGISTER(softmax_focal_loss_forward)
+    .attr("gamma")
+    .attr("alpha")
+    .input(3)
+    .output(1)
+    .apply(softmax_focal_loss_forward_cuda_parrots)
+    .done();
+
+PARROTS_EXTENSION_REGISTER(softmax_focal_loss_backward)
+    .attr("gamma")
+    .attr("alpha")
+    .input(3)
+    .output(2)
+    .apply(softmax_focal_loss_backward_cuda_parrots)
+    .done();

mmcv/ops/csrc/parrots/focal_loss_pytorch.h

+#ifndef FOCAL_LOSS_PYTORCH_H
+#define FOCAL_LOSS_PYTORCH_H
+#include <torch/extension.h>
+using namespace at;
+
+void sigmoid_focal_loss_forward_cuda(Tensor input, Tensor target, Tensor weight,
+                                     Tensor output, float gamma, float alpha);
+
+void sigmoid_focal_loss_backward_cuda(Tensor input, Tensor target,
+                                      Tensor weight, Tensor grad_input,
+                                      float gamma, float alpha);
+
+void softmax_focal_loss_forward_cuda(Tensor input, Tensor target, Tensor weight,
+                                     Tensor output, float gamma, float alpha);
+
+void softmax_focal_loss_backward_cuda(Tensor input, Tensor target,
+                                      Tensor weight, Tensor buff,
+                                      Tensor grad_input, float gamma,
+                                      float alpha);
+#endif  // FOCAL_LOSS_PYTORCH_H

mmcv/ops/csrc/parrots/masked_conv2d.cpp

-#include "parrots_cpp_helper.hpp"
+#include "pytorch_cpp_helper.hpp"
 
-void MaskedIm2colForwardCUDAKernelLauncher(
-    const DArrayLite bottom_data, const DArrayLite mask_h_idx,
-    const DArrayLite mask_w_idx, DArrayLite top_data, const int kernel_h,
-    const int kernel_w, const int pad_h, const int pad_w, cudaStream_t stream);
+#ifdef MMCV_WITH_CUDA
+void MaskedIm2colForwardCUDAKernelLauncher(const Tensor bottom_data,
+                                           const Tensor mask_h_idx,
+                                           const Tensor mask_w_idx,
+                                           Tensor top_data, const int kernel_h,
+                                           const int kernel_w, const int pad_h,
+                                           const int pad_w);
 
-void MaskedCol2imForwardCUDAKernelLaucher(const DArrayLite bottom_data,
-                                          const DArrayLite mask_h_idx,
-                                          const DArrayLite mask_w_idx,
-                                          DArrayLite top_data, const int height,
-                                          const int width, const int channels,
-                                          cudaStream_t stream);
+void MaskedCol2imForwardCUDAKernelLauncher(const Tensor bottom_data,
+                                           const Tensor mask_h_idx,
+                                           const Tensor mask_w_idx,
+                                           Tensor top_data, const int height,
+                                           const int width, const int channels);
 
-void masked_im2col_forward_cuda(CudaContext& ctx, const SSElement& attr,
-                                const OperatorBase::in_list_t& ins,
-                                OperatorBase::out_list_t& outs) {
+void masked_im2col_forward_cuda(const Tensor im, const Tensor mask_h_idx,
+                                const Tensor mask_w_idx, Tensor col,
+                                const int kernel_h, const int kernel_w,
+                                const int pad_h, const int pad_w) {
   // im: (n, ic, h, w), kernel size (kh, kw)
   // kernel: (oc, ic * kh * kw), col: (kh * kw * ic, ow * oh)
-  int kernel_h, kernel_w, pad_h, pad_w;
-  SSAttrs(attr)
-      .get<int>("kernel_h", kernel_h)
-      .get<int>("kernel_w", kernel_w)
-      .get<int>("pad_h", pad_h)
-      .get<int>("pad_w", pad_w)
-      .done();
-
-  const auto& im = ins[0];
-  const auto& mask_h_idx = ins[1];
-  const auto& mask_w_idx = ins[2];
-
-  auto& col = outs[0];
-
-  cudaStream_t stream = getStreamNative<CudaDevice>(ctx.getStream());
   MaskedIm2colForwardCUDAKernelLauncher(im, mask_h_idx, mask_w_idx, col,
-                                        kernel_h, kernel_w, pad_h, pad_w,
-                                        stream);
+                                        kernel_h, kernel_w, pad_h, pad_w);
 }
 
-void masked_col2im_forward_cuda(CudaContext& ctx, const SSElement& attr,
-                                const OperatorBase::in_list_t& ins,
-                                OperatorBase::out_list_t& outs) {
+void masked_col2im_forward_cuda(const Tensor col, const Tensor mask_h_idx,
+                                const Tensor mask_w_idx, Tensor im, int height,
+                                int width, int channels) {
   // im: (n, ic, h, w), kernel size (kh, kw)
   // kernel: (oc, ic * kh * kh), col: (kh * kw * ic, ow * oh)
-  int height, width, channels;
-  SSAttrs(attr)
-      .get<int>("height", height)
-      .get<int>("width", width)
-      .get<int>("channels", channels)
-      .done();
-
-  const auto& col = ins[0];
-  const auto& mask_h_idx = ins[1];
-  const auto& mask_w_idx = ins[2];
-
-  auto& im = outs[0];
-
-  cudaStream_t stream = getStreamNative<CudaDevice>(ctx.getStream());
-  MaskedCol2imForwardCUDAKernelLaucher(col, mask_h_idx, mask_w_idx, im, height,
-                                       width, channels, stream);
+  MaskedCol2imForwardCUDAKernelLauncher(col, mask_h_idx, mask_w_idx, im, height,
+                                        width, channels);
 }
+#endif
 
-PARROTS_EXTENSION_REGISTER(masked_im2col_forward)
-    .attr("kernel_h")
-    .attr("kernel_w")
-    .attr("pad_h")
-    .attr("pad_w")
-    .input(3)
-    .output(1)
-    .apply(masked_im2col_forward_cuda)
-    .done();
+void masked_im2col_forward(const Tensor im, const Tensor mask_h_idx,
+                           const Tensor mask_w_idx, Tensor col,
+                           const int kernel_h, const int kernel_w,
+                           const int pad_h, const int pad_w) {
+  if (im.device().is_cuda()) {
+#ifdef MMCV_WITH_CUDA
+    CHECK_CUDA_INPUT(im);
+    CHECK_CUDA_INPUT(mask_h_idx);
+    CHECK_CUDA_INPUT(mask_w_idx);
+    CHECK_CUDA_INPUT(col);
+    masked_im2col_forward_cuda(im, mask_h_idx, mask_w_idx, col, kernel_h,
+                               kernel_w, pad_h, pad_w);
+#else
+    AT_ERROR("MaskConv is not compiled with GPU support");
+#endif
+  } else {
+    AT_ERROR("MaskConv is not implemented on CPU");
+  }
+}
 
-PARROTS_EXTENSION_REGISTER(masked_col2im_forward)
-    .attr("height")
-    .attr("width")
-    .attr("channels")
-    .input(3)
-    .output(1)
-    .apply(masked_col2im_forward_cuda)
-    .done();
+void masked_col2im_forward(const Tensor col, const Tensor mask_h_idx,
+                           const Tensor mask_w_idx, Tensor im, int height,
+                           int width, int channels) {
+  if (col.device().is_cuda()) {
+#ifdef MMCV_WITH_CUDA
+    CHECK_CUDA_INPUT(col);
+    CHECK_CUDA_INPUT(mask_h_idx);
+    CHECK_CUDA_INPUT(mask_w_idx);
+    CHECK_CUDA_INPUT(im);
+    masked_col2im_forward_cuda(col, mask_h_idx, mask_w_idx, im, height, width,
+                               channels);
+#else
+    AT_ERROR("MaskConv is not compiled with GPU support");
+#endif
+  } else {
+    AT_ERROR("MaskConv is not implemented on CPU");
+  }
+}

mmcv/ops/csrc/parrots/masked_conv2d_cuda.cu

 #include "masked_conv2d_cuda_kernel.cuh"
-#include "parrots_cuda_helper.hpp"
+#include "pytorch_cuda_helper.hpp"
 
-void MaskedIm2colForwardCUDAKernelLauncher(
-    const DArrayLite bottom_data, const DArrayLite mask_h_idx,
-    const DArrayLite mask_w_idx, DArrayLite top_data, const int kernel_h,
-    const int kernel_w, const int pad_h, const int pad_w, cudaStream_t stream) {
-  int channels = bottom_data.dim(1);
-  int height = bottom_data.dim(2);
-  int width = bottom_data.dim(3);
-  int mask_cnt = mask_h_idx.dim(0);
+void MaskedIm2colForwardCUDAKernelLauncher(const Tensor bottom_data,
+                                           const Tensor mask_h_idx,
+                                           const Tensor mask_w_idx,
+                                           Tensor top_data, const int kernel_h,
+                                           const int kernel_w, const int pad_h,
+                                           const int pad_w) {
+  int channels = bottom_data.size(1);
+  int height = bottom_data.size(2);
+  int width = bottom_data.size(3);
+  int mask_cnt = mask_h_idx.size(0);
   int output_size = mask_cnt * channels;
 
-  PARROTS_DISPATCH_FLOATING_TYPES_AND_HALF(
-      bottom_data.elemType().prim(), ([&] {
+  at::cuda::CUDAGuard device_guard(bottom_data.device());
+  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+      bottom_data.scalar_type(), "MaskedIm2colLaucherForward", ([&] {
+        const scalar_t *bottom_data_ = bottom_data.data_ptr<scalar_t>();
+        const int64_t *mask_h_idx_ = mask_h_idx.data_ptr<int64_t>();
+        const int64_t *mask_w_idx_ = mask_w_idx.data_ptr<int64_t>();
+        scalar_t *top_data_ = top_data.data_ptr<scalar_t>();
         MaskedIm2colForward<scalar_t>
             <<<GET_BLOCKS(output_size), THREADS_PER_BLOCK, 0, stream>>>(
-                output_size, bottom_data.ptr<scalar_t>(), height, width,
-                kernel_h, kernel_w, pad_h, pad_w, mask_h_idx.ptr<int64_t>(),
-                mask_w_idx.ptr<int64_t>(), mask_cnt, top_data.ptr<scalar_t>());
+                output_size, bottom_data_, height, width, kernel_h, kernel_w,
+                pad_h, pad_w, mask_h_idx_, mask_w_idx_, mask_cnt, top_data_);
       }));
-
-  PARROTS_CUDA_CHECK(cudaGetLastError());
+  AT_CUDA_CHECK(cudaGetLastError());
 }
 
-void MaskedCol2imForwardCUDAKernelLaucher(const DArrayLite bottom_data,
-                                          const DArrayLite mask_h_idx,
-                                          const DArrayLite mask_w_idx,
-                                          DArrayLite top_data, const int height,
-                                          const int width, const int channels,
-                                          cudaStream_t stream) {
-  int mask_cnt = mask_h_idx.dim(0);
+void MaskedCol2imForwardCUDAKernelLauncher(
+    const Tensor bottom_data, const Tensor mask_h_idx, const Tensor mask_w_idx,
+    Tensor top_data, const int height, const int width, const int channels) {
+  int mask_cnt = mask_h_idx.size(0);
   int output_size = mask_cnt * channels;
 
-  PARROTS_DISPATCH_FLOATING_TYPES_AND_HALF(
-      bottom_data.elemType().prim(), ([&] {
+  at::cuda::CUDAGuard device_guard(bottom_data.device());
+  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+      bottom_data.scalar_type(), "MaskedCol2imLaucherForward", ([&] {
+        const scalar_t *bottom_data_ = bottom_data.data_ptr<scalar_t>();
+        const int64_t *mask_h_idx_ = mask_h_idx.data_ptr<int64_t>();
+        const int64_t *mask_w_idx_ = mask_w_idx.data_ptr<int64_t>();
+        scalar_t *top_data_ = top_data.data_ptr<scalar_t>();
+
         MaskedCol2imForward<scalar_t>
             <<<GET_BLOCKS(output_size), THREADS_PER_BLOCK, 0, stream>>>(
-                output_size, bottom_data.ptr<scalar_t>(), height, width,
-                channels, mask_h_idx.ptr<int64_t>(), mask_w_idx.ptr<int64_t>(),
-                mask_cnt, top_data.ptr<scalar_t>());
+                output_size, bottom_data_, height, width, channels, mask_h_idx_,
+                mask_w_idx_, mask_cnt, top_data_);
       }));
-
-  PARROTS_CUDA_CHECK(cudaGetLastError());
+  AT_CUDA_CHECK(cudaGetLastError());
 }

mmcv/ops/csrc/parrots/masked_conv2d_parrots.cpp

+#include <parrots/compute/aten.hpp>
+#include <parrots/extension.hpp>
+#include <parrots/foundation/ssattrs.hpp>
+
+#include "masked_conv2d_pytorch.h"
+
+using namespace parrots;
+
+void masked_im2col_forward_cuda_parrots(CudaContext& ctx, const SSElement& attr,
+                                        const OperatorBase::in_list_t& ins,
+                                        OperatorBase::out_list_t& outs) {
+  // im: (n, ic, h, w), kernel size (kh, kw)
+  // kernel: (oc, ic * kh * kw), col: (kh * kw * ic, ow * oh)
+  int kernel_h, kernel_w, pad_h, pad_w;
+  SSAttrs(attr)
+      .get<int>("kernel_h", kernel_h)
+      .get<int>("kernel_w", kernel_w)
+      .get<int>("pad_h", pad_h)
+      .get<int>("pad_w", pad_w)
+      .done();
+
+  const auto& im = buildATensor(ctx, ins[0]);
+  const auto& mask_h_idx = buildATensor(ctx, ins[1]);
+  const auto& mask_w_idx = buildATensor(ctx, ins[2]);
+
+  auto col = buildATensor(ctx, outs[0]);
+  masked_im2col_forward_cuda(im, mask_h_idx, mask_w_idx, col, kernel_h,
+                             kernel_w, pad_h, pad_w);
+}
+
+void masked_col2im_forward_cuda_parrots(CudaContext& ctx, const SSElement& attr,
+                                        const OperatorBase::in_list_t& ins,
+                                        OperatorBase::out_list_t& outs) {
+  // im: (n, ic, h, w), kernel size (kh, kw)
+  // kernel: (oc, ic * kh * kh), col: (kh * kw * ic, ow * oh)
+  int height, width, channels;
+  SSAttrs(attr)
+      .get<int>("height", height)
+      .get<int>("width", width)
+      .get<int>("channels", channels)
+      .done();
+
+  const auto& col = buildATensor(ctx, ins[0]);
+  const auto& mask_h_idx = buildATensor(ctx, ins[1]);
+  const auto& mask_w_idx = buildATensor(ctx, ins[2]);
+
+  auto im = buildATensor(ctx, outs[0]);
+  masked_col2im_forward_cuda(col, mask_h_idx, mask_w_idx, im, height, width,
+                             channels);
+}
+
+PARROTS_EXTENSION_REGISTER(masked_im2col_forward)
+    .attr("kernel_h")
+    .attr("kernel_w")
+    .attr("pad_h")
+    .attr("pad_w")
+    .input(3)
+    .output(1)
+    .apply(masked_im2col_forward_cuda_parrots)
+    .done();
+
+PARROTS_EXTENSION_REGISTER(masked_col2im_forward)
+    .attr("height")
+    .attr("width")
+    .attr("channels")
+    .input(3)
+    .output(1)
+    .apply(masked_col2im_forward_cuda_parrots)
+    .done();