Add Torch Selective macros in all C++ Ops for better support on mobile (#3218)

* Adding TORCH_SELECTIVE_* macros on op registration.

* Adding torchvision namespace.

torchvision/csrc/ops/cuda/ps_roi_pool_kernel.cu

@@ -276,8 +276,12 @@ at::Tensor ps_roi_pool_backward_kernel(
 } // namespace
 
 TORCH_LIBRARY_IMPL(torchvision, CUDA, m) {
-  m.impl("ps_roi_pool", ps_roi_pool_forward_kernel);
-  m.impl("_ps_roi_pool_backward", ps_roi_pool_backward_kernel);
+  m.impl(
+      TORCH_SELECTIVE_NAME("torchvision::ps_roi_pool"),
+      TORCH_FN(ps_roi_pool_forward_kernel));
+  m.impl(
+      TORCH_SELECTIVE_NAME("torchvision::_ps_roi_pool_backward"),
+      TORCH_FN(ps_roi_pool_backward_kernel));
 }
 
 } // namespace ops

torchvision/csrc/ops/cuda/roi_align_kernel.cu

@@ -449,8 +449,12 @@ at::Tensor roi_align_backward_kernel(
 } // namespace
 
 TORCH_LIBRARY_IMPL(torchvision, CUDA, m) {
-  m.impl("roi_align", roi_align_forward_kernel);
-  m.impl("_roi_align_backward", roi_align_backward_kernel);
+  m.impl(
+      TORCH_SELECTIVE_NAME("torchvision::roi_align"),
+      TORCH_FN(roi_align_forward_kernel));
+  m.impl(
+      TORCH_SELECTIVE_NAME("torchvision::_roi_align_backward"),
+      TORCH_FN(roi_align_backward_kernel));
 }
 
 } // namespace ops

torchvision/csrc/ops/cuda/roi_pool_kernel.cu

@@ -260,8 +260,12 @@ at::Tensor roi_pool_backward_kernel(
 } // namespace
 
 TORCH_LIBRARY_IMPL(torchvision, CUDA, m) {
-  m.impl("roi_pool", roi_pool_forward_kernel);
-  m.impl("_roi_pool_backward", roi_pool_backward_kernel);
+  m.impl(
+      TORCH_SELECTIVE_NAME("torchvision::roi_pool"),
+      TORCH_FN(roi_pool_forward_kernel));
+  m.impl(
+      TORCH_SELECTIVE_NAME("torchvision::_roi_pool_backward"),
+      TORCH_FN(roi_pool_backward_kernel));
 }
 
 } // namespace ops

torchvision/csrc/ops/deform_conv2d.cpp

@@ -84,10 +84,10 @@ _deform_conv2d_backward(
 } // namespace detail
 
 TORCH_LIBRARY_FRAGMENT(torchvision, m) {
-  m.def(
-      "deform_conv2d(Tensor input, Tensor weight, Tensor offset, Tensor mask, Tensor bias, int stride_h, int stride_w, int pad_h, int pad_w, int dilation_h, int dilation_w, int groups, int offset_groups, bool use_mask) -> Tensor");
-  m.def(
-      "_deform_conv2d_backward(Tensor grad, Tensor input, Tensor weight, Tensor offset, Tensor mask, Tensor bias, int stride_h, int stride_w, int pad_h, int pad_w, int dilation_h, int dilation_w, int groups, int offset_groups, bool use_mask) -> (Tensor, Tensor, Tensor, Tensor, Tensor)");
+  m.def(TORCH_SELECTIVE_SCHEMA(
+      "torchvision::deform_conv2d(Tensor input, Tensor weight, Tensor offset, Tensor mask, Tensor bias, int stride_h, int stride_w, int pad_h, int pad_w, int dilation_h, int dilation_w, int groups, int offset_groups, bool use_mask) -> Tensor"));
+  m.def(TORCH_SELECTIVE_SCHEMA(
+      "torchvision::_deform_conv2d_backward(Tensor grad, Tensor input, Tensor weight, Tensor offset, Tensor mask, Tensor bias, int stride_h, int stride_w, int pad_h, int pad_w, int dilation_h, int dilation_w, int groups, int offset_groups, bool use_mask) -> (Tensor, Tensor, Tensor, Tensor, Tensor)"));
 }
 
 } // namespace ops

torchvision/csrc/ops/nms.cpp

@@ -16,7 +16,8 @@ at::Tensor nms(
 }
 
 TORCH_LIBRARY_FRAGMENT(torchvision, m) {
-  m.def("nms(Tensor dets, Tensor scores, float iou_threshold) -> Tensor");
+  m.def(TORCH_SELECTIVE_SCHEMA(
+      "torchvision::nms(Tensor dets, Tensor scores, float iou_threshold) -> Tensor"));
 }
 
 } // namespace ops

torchvision/csrc/ops/ps_roi_align.cpp

@@ -54,10 +54,10 @@ at::Tensor _ps_roi_align_backward(
 } // namespace detail
 
 TORCH_LIBRARY_FRAGMENT(torchvision, m) {
-  m.def(
-      "ps_roi_align(Tensor input, Tensor rois, float spatial_scale, int pooled_height, int pooled_width, int sampling_ratio) -> (Tensor, Tensor)");
-  m.def(
-      "_ps_roi_align_backward(Tensor grad, Tensor rois, Tensor channel_mapping, float spatial_scale, int pooled_height, int pooled_width, int sampling_ratio, int batch_size, int channels, int height, int width) -> Tensor");
+  m.def(TORCH_SELECTIVE_SCHEMA(
+      "torchvision::ps_roi_align(Tensor input, Tensor rois, float spatial_scale, int pooled_height, int pooled_width, int sampling_ratio) -> (Tensor, Tensor)"));
+  m.def(TORCH_SELECTIVE_SCHEMA(
+      "torchvision::_ps_roi_align_backward(Tensor grad, Tensor rois, Tensor channel_mapping, float spatial_scale, int pooled_height, int pooled_width, int sampling_ratio, int batch_size, int channels, int height, int width) -> Tensor"));
 }
 
 } // namespace ops

torchvision/csrc/ops/ps_roi_pool.cpp

@@ -50,10 +50,10 @@ at::Tensor _ps_roi_pool_backward(
 } // namespace detail
 
 TORCH_LIBRARY_FRAGMENT(torchvision, m) {
-  m.def(
-      "ps_roi_pool(Tensor input, Tensor rois, float spatial_scale, int pooled_height, int pooled_width) -> (Tensor, Tensor)");
-  m.def(
-      "_ps_roi_pool_backward(Tensor grad, Tensor rois, Tensor channel_mapping, float spatial_scale, int pooled_height, int pooled_width, int batch_size, int channels, int height, int width) -> Tensor");
+  m.def(TORCH_SELECTIVE_SCHEMA(
+      "torchvision::ps_roi_pool(Tensor input, Tensor rois, float spatial_scale, int pooled_height, int pooled_width) -> (Tensor, Tensor)"));
+  m.def(TORCH_SELECTIVE_SCHEMA(
+      "torchvision::_ps_roi_pool_backward(Tensor grad, Tensor rois, Tensor channel_mapping, float spatial_scale, int pooled_height, int pooled_width, int batch_size, int channels, int height, int width) -> Tensor"));
 }
 
 } // namespace ops

torchvision/csrc/ops/roi_align.cpp

@@ -64,10 +64,10 @@ at::Tensor _roi_align_backward(
 } // namespace detail
 
 TORCH_LIBRARY_FRAGMENT(torchvision, m) {
-  m.def(
-      "roi_align(Tensor input, Tensor rois, float spatial_scale, int pooled_height, int pooled_width, int sampling_ratio, bool aligned) -> Tensor");
-  m.def(
-      "_roi_align_backward(Tensor grad, Tensor rois, float spatial_scale, int pooled_height, int pooled_width, int batch_size, int channels, int height, int width, int sampling_ratio, bool aligned) -> Tensor");
+  m.def(TORCH_SELECTIVE_SCHEMA(
+      "torchvision::roi_align(Tensor input, Tensor rois, float spatial_scale, int pooled_height, int pooled_width, int sampling_ratio, bool aligned) -> Tensor"));
+  m.def(TORCH_SELECTIVE_SCHEMA(
+      "torchvision::_roi_align_backward(Tensor grad, Tensor rois, float spatial_scale, int pooled_height, int pooled_width, int batch_size, int channels, int height, int width, int sampling_ratio, bool aligned) -> Tensor"));
 }
 
 } // namespace ops

torchvision/csrc/ops/roi_pool.cpp

@@ -49,10 +49,10 @@ at::Tensor _roi_pool_backward(
 } // namespace detail
 
 TORCH_LIBRARY_FRAGMENT(torchvision, m) {
-  m.def(
-      "roi_pool(Tensor input, Tensor rois, float spatial_scale, int pooled_height, int pooled_width) -> (Tensor, Tensor)");
-  m.def(
-      "_roi_pool_backward(Tensor grad, Tensor rois, Tensor argmax, float spatial_scale, int pooled_height, int pooled_width, int batch_size, int channels, int height, int width) -> Tensor");
+  m.def(TORCH_SELECTIVE_SCHEMA(
+      "torchvision::roi_pool(Tensor input, Tensor rois, float spatial_scale, int pooled_height, int pooled_width) -> (Tensor, Tensor)"));
+  m.def(TORCH_SELECTIVE_SCHEMA(
+      "torchvision::_roi_pool_backward(Tensor grad, Tensor rois, Tensor argmax, float spatial_scale, int pooled_height, int pooled_width, int batch_size, int channels, int height, int width) -> Tensor"));
 }
 
 } // namespace ops