Create foreach cpp file a separate header file with "public" functions.

Author: datumbox

torchvision/csrc/cpu/nms_kernel.cpp

@@ -1,4 +1,4 @@
-#include "vision_cpu.h"
+#include "nms_kernel.h"
 
 namespace {
 

torchvision/csrc/cpu/nms_kernel.h

@@ -0,0 +1,9 @@
+#pragma once
+
+#include <ATen/ATen.h>
+#include "../macros.h"
+
+VISION_API at::Tensor nms_cpu(
+    const at::Tensor& dets,
+    const at::Tensor& scores,
+    double iou_threshold);

torchvision/csrc/cpu/vision_cpu.h

@@ -4,11 +4,6 @@
 
 // TODO: Delete this file once all the methods are gone
 
-VISION_API at::Tensor nms_cpu(
-    const at::Tensor& dets,
-    const at::Tensor& scores,
-    double iou_threshold);
-
 VISION_API std::tuple<at::Tensor, at::Tensor> PSROIAlign_forward_cpu(
     const at::Tensor& input,
     const at::Tensor& rois,

torchvision/csrc/cuda/nms_kernel.cu

@@ -3,6 +3,7 @@
 #include <c10/cuda/CUDAGuard.h>
 
 #include "cuda_helpers.h"
+#include "nms_kernel.h"
 
 #include <iostream>
 #include <vector>

torchvision/csrc/cuda/nms_kernel.h

@@ -0,0 +1,9 @@
+#pragma once
+
+#include <ATen/ATen.h>
+#include "../macros.h"
+
+VISION_API at::Tensor nms_cuda(
+    const at::Tensor& dets,
+    const at::Tensor& scores,
+    double iou_threshold);

torchvision/csrc/cuda/vision_cuda.h

@@ -4,11 +4,6 @@
 
 // TODO: Delete this file once all the methods are gone
 
-VISION_API at::Tensor nms_cuda(
-    const at::Tensor& dets,
-    const at::Tensor& scores,
-    double iou_threshold);
-
 VISION_API std::tuple<at::Tensor, at::Tensor> PSROIAlign_forward_cuda(
     const at::Tensor& input,
     const at::Tensor& rois,

torchvision/csrc/nms.cpp

@@ -1,17 +1,10 @@
-#pragma once
+#include "nms.h"
+#include <torch/extension.h>
 
-#include "cpu/vision_cpu.h"
-
-#ifdef WITH_CUDA
-#include "autocast.h"
-#include "cuda/vision_cuda.h"
-#endif
-#ifdef WITH_HIP
-#include "autocast.h"
-#include "hip/vision_cuda.h"
+#if defined(WITH_CUDA) || defined(WITH_HIP)
+#include <ATen/autocast_mode.h>
 #endif
 
-// nms dispatch nexus
 at::Tensor nms(
     const at::Tensor& dets,
     const at::Tensor& scores,

torchvision/csrc/nms.h

@@ -0,0 +1,24 @@
+#pragma once
+
+#include "cpu/nms_kernel.h"
+
+#ifdef WITH_CUDA
+#include "cuda/nms_kernel.h"
+#endif
+#ifdef WITH_HIP
+#include "hip/nms_kernel.h"
+#endif
+
+// C++ Forward
+at::Tensor nms(
+    const at::Tensor& dets,
+    const at::Tensor& scores,
+    double iou_threshold);
+
+// Autocast Forward
+#if defined(WITH_CUDA) || defined(WITH_HIP)
+at::Tensor nms_autocast(
+    const at::Tensor& dets,
+    const at::Tensor& scores,
+    double iou_threshold);
+#endif