Add RoI torchvision ops

src/ATen/native/xpu/PsRoiAlign.cpp

@@ -0,0 +1,69 @@
+#include <ATen/core/Tensor.h>
+#include <ATen/native/xpu/sycl/PsRoiAlignKernels.h>
+#include <comm/XPUGuard.h>
+#include <comm/xpu_aten.h>
+#include <torch/library.h>
+namespace at::native::xpu {
+
+std::tuple<at::Tensor, at::Tensor> ps_roi_align(
+    const at::Tensor& input,
+    const at::Tensor& rois,
+    double spatial_scale,
+    int64_t pooled_height,
+    int64_t pooled_width,
+    int64_t sampling_ratio) {
+  TORCH_CHECK(input.is_xpu(), "input must be a XPU tensor");
+  TORCH_CHECK(rois.is_xpu(), "rois must be a XPU tensor");
+  TORCH_CHECK(rois.size(1) == 5, "rois must have shape as Tensor[K, 5]");
+
+  at::TensorArg input_t{input, "input", 1}, rois_t{rois, "rois", 2};
+
+  at::CheckedFrom c = "ps_roi_align_forward_kernel";
+  at::checkAllSameGPU(c, {input_t, rois_t});
+  at::checkAllSameType(c, {input_t, rois_t});
+
+  c10::DeviceGuard device_guard(input.device());
+  return ps_roi_align_kernel(
+      input, rois, spatial_scale, pooled_height, pooled_width, sampling_ratio);
+}
+
+at::Tensor _ps_roi_align_backward(
+    const at::Tensor& grad,
+    const at::Tensor& rois,
+    const at::Tensor& channel_mapping,
+    double spatial_scale,
+    int64_t pooled_height,
+    int64_t pooled_width,
+    int64_t batch_size,
+    int64_t channels,
+    int64_t height,
+    int64_t width,
+    int64_t sampling_ratio) {
+  TORCH_CHECK(grad.is_xpu(), "grad must be a XPU tensor");
+  TORCH_CHECK(rois.is_xpu(), "rois must be a XPU tensor");
+  TORCH_CHECK(channel_mapping.is_xpu(), "channel_mapping must be a XPU tensor");
+
+  at::TensorArg grad_t{grad, "grad", 1}, rois_t{rois, "rois", 2},
+      channel_mapping_t{channel_mapping, "channel_mapping", 3};
+
+  at::CheckedFrom c = "ps_roi_align_backward_kernel";
+  at::checkAllSameGPU(c, {grad_t, rois_t, channel_mapping_t});
+  at::checkAllSameType(c, {grad_t, rois_t});
+
+  c10::DeviceGuard device_guard(grad.device());
+
+  return ps_roi_align_backward_kernel(
+      grad,
+      rois,
+      channel_mapping,
+      spatial_scale,
+      pooled_height,
+      pooled_width,
+      batch_size,
+      channels,
+      height,
+      width,
+      sampling_ratio);
+}
+
+} // namespace at::native::xpu

src/ATen/native/xpu/PsRoiPool.cpp

@@ -0,0 +1,66 @@
+#include <ATen/core/Tensor.h>
+#include <ATen/native/xpu/sycl/PsRoiPoolKernels.h>
+#include <comm/XPUGuard.h>
+#include <comm/xpu_aten.h>
+#include <torch/library.h>
+namespace at::native::xpu {
+
+std::tuple<at::Tensor, at::Tensor> ps_roi_pool(
+    const at::Tensor& input,
+    const at::Tensor& rois,
+    double spatial_scale,
+    int64_t pooled_height,
+    int64_t pooled_width) {
+  TORCH_CHECK(input.is_xpu(), "input must be a XPU tensor");
+  TORCH_CHECK(rois.is_xpu(), "rois must be a XPU tensor");
+  TORCH_CHECK(rois.size(1) == 5, "rois must have shape as Tensor[K, 5]");
+
+  at::TensorArg input_t{input, "input", 1}, rois_t{rois, "rois", 2};
+
+  at::CheckedFrom c = "ps_roi_pool_forward_kernel";
+  at::checkAllSameGPU(c, {input_t, rois_t});
+  at::checkAllSameType(c, {input_t, rois_t});
+
+  c10::DeviceGuard device_guard(input.device());
+  return ps_roi_pool_kernel(
+      input, rois, spatial_scale, pooled_height, pooled_width);
+}
+
+at::Tensor _ps_roi_pool_backward(
+    const at::Tensor& grad,
+    const at::Tensor& rois,
+    const at::Tensor& channel_mapping,
+    double spatial_scale,
+    int64_t pooled_height,
+    int64_t pooled_width,
+    int64_t batch_size,
+    int64_t channels,
+    int64_t height,
+    int64_t width) {
+  TORCH_CHECK(grad.is_xpu(), "grad must be a XPU tensor");
+  TORCH_CHECK(rois.is_xpu(), "rois must be a XPU tensor");
+  TORCH_CHECK(channel_mapping.is_xpu(), "channel_mapping must be a XPU tensor");
+
+  at::TensorArg grad_t{grad, "grad", 1}, rois_t{rois, "rois", 2},
+      channel_mapping_t{channel_mapping, "channel_mapping", 3};
+
+  at::CheckedFrom c = "ps_roi_pool_backward_kernel";
+  at::checkAllSameGPU(c, {grad_t, rois_t, channel_mapping_t});
+  at::checkAllSameType(c, {grad_t, rois_t});
+
+  c10::DeviceGuard device_guard(grad.device());
+
+  return ps_roi_pool_backward_kernel(
+      grad,
+      rois,
+      channel_mapping,
+      spatial_scale,
+      pooled_height,
+      pooled_width,
+      batch_size,
+      channels,
+      height,
+      width);
+}
+
+} // namespace at::native::xpu

src/ATen/native/xpu/RoiPool.cpp

@@ -0,0 +1,66 @@
+#include <ATen/core/Tensor.h>
+#include <ATen/native/xpu/sycl/RoiPoolKernels.h>
+#include <comm/XPUGuard.h>
+#include <comm/xpu_aten.h>
+#include <torch/library.h>
+namespace at::native::xpu {
+
+std::tuple<at::Tensor, at::Tensor> roi_pool(
+    const at::Tensor& input,
+    const at::Tensor& rois,
+    double spatial_scale,
+    int64_t pooled_height,
+    int64_t pooled_width) {
+  TORCH_CHECK(input.is_xpu(), "input must be a XPU tensor");
+  TORCH_CHECK(rois.is_xpu(), "rois must be a XPU tensor");
+  TORCH_CHECK(rois.size(1) == 5, "rois must have shape as Tensor[K, 5]");
+
+  at::TensorArg input_t{input, "input", 1}, rois_t{rois, "rois", 2};
+
+  at::CheckedFrom c = "roi_pool_forward_kernel";
+  at::checkAllSameGPU(c, {input_t, rois_t});
+  at::checkAllSameType(c, {input_t, rois_t});
+
+  c10::DeviceGuard device_guard(input.device());
+  return roi_pool_kernel(
+      input, rois, spatial_scale, pooled_height, pooled_width);
+}
+
+at::Tensor _roi_pool_backward(
+    const at::Tensor& grad,
+    const at::Tensor& rois,
+    const at::Tensor& argmax,
+    double spatial_scale,
+    int64_t pooled_height,
+    int64_t pooled_width,
+    int64_t batch_size,
+    int64_t channels,
+    int64_t height,
+    int64_t width) {
+  TORCH_CHECK(grad.is_xpu(), "grad must be a XPU tensor");
+  TORCH_CHECK(rois.is_xpu(), "rois must be a XPU tensor");
+  TORCH_CHECK(argmax.is_xpu(), "argmax must be a XPU tensor");
+
+  at::TensorArg grad_t{grad, "grad", 1}, rois_t{rois, "rois", 2},
+      argmax_t{argmax, "argmax", 3};
+
+  at::CheckedFrom c = "roi_pool_backward_kernel";
+  at::checkAllSameGPU(c, {grad_t, rois_t, argmax_t});
+  at::checkAllSameType(c, {grad_t, rois_t});
+
+  c10::DeviceGuard device_guard(grad.device());
+
+  return roi_pool_backward_kernel(
+      grad,
+      rois,
+      argmax,
+      spatial_scale,
+      pooled_height,
+      pooled_width,
+      batch_size,
+      channels,
+      height,
+      width);
+}
+
+} // namespace at::native::xpu

src/ATen/native/xpu/XPUFallback.template

@@ -26,6 +26,12 @@ namespace native::xpu {
 Tensor nms(const Tensor& dets, const Tensor& scores, double iou_threshold_);
 Tensor roi_align(const Tensor& input, const Tensor& rois, double spatial_scale, int64_t pooled_height, int64_t pooled_width, int64_t sampling_ratio, bool aligned);
 Tensor _roi_align_backward(const Tensor& grad, const Tensor& rois, double spatial_scale, int64_t pooled_height, int64_t pooled_width, int64_t batch_size, int64_t channels, int64_t height, int64_t width, int64_t sampling_ratio, bool aligned);
+std::tuple<Tensor, Tensor> ps_roi_align(const Tensor& input, const Tensor& rois, double spatial_scale, int64_t pooled_height, int64_t pooled_width, int64_t sampling_ratio);
+Tensor _ps_roi_align_backward(const Tensor& grad, const Tensor& rois, const Tensor& channel_mapping, double spatial_scale, int64_t pooled_height, int64_t pooled_width, int64_t sampling_ratio, int64_t batch_size, int64_t channels, int64_t height, int64_t width);
+std::tuple<Tensor, Tensor> roi_pool(const Tensor& input, const Tensor& rois, double spatial_scale, int64_t pooled_height, int64_t pooled_width);
+Tensor _roi_pool_backward(const Tensor& grad, const Tensor& rois, const Tensor& argmax, double spatial_scale, int64_t pooled_height, int64_t pooled_width, int64_t batch_size, int64_t channels, int64_t height, int64_t width);
+std::tuple<Tensor, Tensor> ps_roi_pool(const Tensor& input, const Tensor& rois, double spatial_scale, int64_t pooled_height, int64_t pooled_width);
+Tensor _ps_roi_pool_backward(const Tensor& grad, const Tensor& rois, const Tensor& channel_mapping, double spatial_scale, int64_t pooled_height, int64_t pooled_width, int64_t batch_size, int64_t channels, int64_t height, int64_t width);
 }
 
 // Register op's implementation lazily since sometimes the op is not defined,
@@ -38,6 +44,12 @@ static std::map<std::string, bool> torchvision_ops_dispatching_table_ = {
   {"torchvision::nms", false},
   {"torchvision::roi_align", false},
   {"torchvision::_roi_align_backward", false},
+  {"torchvision::ps_roi_align", false},
+  {"torchvision::_ps_roi_align_backward", false},
+  {"torchvision::roi_pool", false},
+  {"torchvision::_roi_pool_backward", false},
+  {"torchvision::ps_roi_pool", false},
+  {"torchvision::_ps_roi_pool_backward", false},
 };
 
 // Return:
@@ -56,6 +68,18 @@ static bool lazy_registration_and_redispatch(
             TORCH_SELECTIVE_NAME("torchvision::roi_align"),TORCH_FN(at::native::xpu::roi_align));
         m.impl(
             TORCH_SELECTIVE_NAME("torchvision::_roi_align_backward"),TORCH_FN(at::native::xpu::_roi_align_backward));
+        m.impl(
+            TORCH_SELECTIVE_NAME("torchvision::ps_roi_align"),TORCH_FN(at::native::xpu::ps_roi_align));
+        m.impl(
+            TORCH_SELECTIVE_NAME("torchvision::_ps_roi_align_backward"),TORCH_FN(at::native::xpu::_ps_roi_align_backward));
+        m.impl(
+            TORCH_SELECTIVE_NAME("torchvision::roi_pool"),TORCH_FN(at::native::xpu::roi_pool));
+        m.impl(
+            TORCH_SELECTIVE_NAME("torchvision::_roi_pool_backward"),TORCH_FN(at::native::xpu::_roi_pool_backward));
+        m.impl(
+            TORCH_SELECTIVE_NAME("torchvision::ps_roi_pool"),TORCH_FN(at::native::xpu::ps_roi_pool));
+        m.impl(
+            TORCH_SELECTIVE_NAME("torchvision::_ps_roi_pool_backward"),TORCH_FN(at::native::xpu::_ps_roi_pool_backward));
       };
 
   static const torch::detail::TorchLibraryInit