[vulkan] Ops registration to TORCH_LIBRARY_IMPL (pytorch#42194)

Summary: Pull Request resolved: pytorch#42194 Test Plan: Imported from OSS Reviewed By: AshkanAliabadi Differential Revision: D22803036 Pulled By: IvanKobzarev fbshipit-source-id: 2f402541aecf887d78f650bf05d758a0e403bc4d
yanbing-j · Aug 7, 2020 · 3c66a37 · 3c66a37
1 parent 4eb02ad
commit 3c66a37
Show file tree

Hide file tree

Showing 16 changed files with 328 additions and 352 deletions.
diff --git a/BUILD.bazel b/BUILD.bazel
@@ -294,8 +294,8 @@ filegroup(
 )
 
 filegroup(
-    name = "aten_native_vulkan_stub",
-    srcs = glob(["aten/src/ATen/native/vulkan/stub/*.cpp"]),
+    name = "aten_base_vulkan",
+    srcs = glob(["aten/src/ATen/vulkan/*.cpp"]),
 )
 
 filegroup(
@@ -646,13 +646,13 @@ cc_library(
         ":ATen_CORE_SRCS",
         ":ATen_QUANTIZED_SRCS",
         ":aten_base_cpp",
+        ":aten_base_vulkan",
         ":aten_native_cpp",
         ":aten_native_mkl_cpp",
         ":aten_native_mkldnn_cpp",
         ":aten_native_quantized_cpp",
         ":aten_native_sparse_cpp",
         ":aten_native_xnnpack",
-        ":aten_native_vulkan_stub",
         ":aten_src_ATen_config",
         ":generated_cpp",
     ],

diff --git a/aten/src/ATen/CMakeLists.txt b/aten/src/ATen/CMakeLists.txt
@@ -63,8 +63,8 @@ file(GLOB mkldnn_cpp "mkldnn/*.cpp")
 file(GLOB native_cpp "native/*.cpp")
 file(GLOB native_mkl_cpp "native/mkl/*.cpp")
 file(GLOB native_mkldnn_cpp "native/mkldnn/*.cpp")
+file(GLOB vulkan_cpp "vulkan/*.cpp")
 file(GLOB native_vulkan_cpp "native/vulkan/*.cpp")
-file(GLOB native_vulkan_stub_cpp "native/vulkan/stub/*.cpp")
 file(GLOB native_sparse_cpp "native/sparse/*.cpp")
 file(GLOB native_quantized_cpp
             "native/quantized/*.cpp"
@@ -109,9 +109,9 @@ if(AT_MKLDNN_ENABLED)
   set(all_cpu_cpp ${all_cpu_cpp} ${mkldnn_cpp})
 endif()
 if(USE_VULKAN)
-  set(all_cpu_cpp ${all_cpu_cpp} ${native_vulkan_cpp} ${vulkan_generated_cpp})
+  set(all_cpu_cpp ${all_cpu_cpp} ${vulkan_cpp} ${native_vulkan_cpp} ${vulkan_generated_cpp})
 else()
-  set(all_cpu_cpp ${all_cpu_cpp} ${native_vulkan_stub_cpp})
+  set(all_cpu_cpp ${all_cpu_cpp} ${vulkan_cpp})
 endif()
 
 if(USE_CUDA AND USE_ROCM)

diff --git a/aten/src/ATen/native/AdaptiveAveragePooling.cpp b/aten/src/ATen/native/AdaptiveAveragePooling.cpp
@@ -2,9 +2,6 @@
 #include <ATen/NativeFunctions.h>
 #include <ATen/Parallel.h>
 #include <tuple>
-#ifdef USE_VULKAN
-#include <ATen/native/vulkan/VulkanAten.h>
-#endif
 
 
 namespace at {
@@ -328,11 +325,6 @@ namespace {
     if (input.is_mkldnn()) {
       return at::mkldnn_adaptive_avg_pool2d(input, output_size);
     }
-#ifdef USE_VULKAN
-    if (input.is_vulkan()) {
-      return at::native::vulkan_adaptive_avg_pool2d(input, output_size);
-    }
-#endif
 
     // TODO: fastpath for Channels_last should be explored later;
     if (input.suggest_memory_format() == at::MemoryFormat::Contiguous && !input.is_quantized() && output_size[0] == 1 && output_size[1] == 1) {

diff --git a/aten/src/ATen/native/Convolution.cpp b/aten/src/ATen/native/Convolution.cpp
@@ -12,9 +12,6 @@
 #if AT_NNPACK_ENABLED()
 #include <nnpack.h>
 #endif
-#ifdef USE_VULKAN
-#include <ATen/native/vulkan/VulkanAten.h>
-#endif
 
 
 constexpr int MIOPEN_DIM_MAX = 5;
@@ -50,7 +47,6 @@ struct ConvParams {
   bool use_mkldnn(const at::Tensor& input, const at::Tensor& weight) const;
   bool use_nnpack(const at::Tensor& input) const;
   bool use_xnnpack(const at::Tensor& input, const at::Tensor& weight, const at::Tensor& bias) const;
-  bool use_vulkan(const at::Tensor& input, const at::Tensor& weight) const;
   bool is_depthwise(const at::Tensor& input, const at::Tensor& weight) const;
 };
 
@@ -278,20 +274,6 @@ auto ConvParams::use_xnnpack(
   return false;
 }
 
-auto ConvParams::use_vulkan(
-        const at::Tensor &input, const at::Tensor& weight) const -> bool {
-#ifdef USE_VULKAN
-  if (!(input.is_vulkan() && input.scalar_type() == kFloat &&
-        !transposed && input.ndimension() == 4)) {
-    return false;
-  }
-  return (groups == 1) || (input.size(1) == groups && groups > 1 &&
-                           weight.size(0) % input.size(1) == 0);
-#else
-  return false;
-#endif
-}
-
 // We currently only have depthwise support for the case where groups ==
 // nInputPlane and nInputPlane == nOutputPlane (the latter due to the lack of
 // a depthwise multiplier)
@@ -690,12 +672,6 @@ at::Tensor _convolution(
         output = at::miopen_depthwise_convolution(
             input.contiguous(), weight, bias,
             padding, stride, dilation, params.groups, params.benchmark, params.deterministic);
-#ifdef USE_VULKAN
-      } else if (params.use_vulkan(input, weight)) {
-        output = at::native::vulkan_convolution(
-            input, weight, bias,
-            params.padding, params.stride, params.dilation, params.groups);
-#endif
       } else {
           output = at::thnn_conv_depthwise2d(input.contiguous(), weight, kernel_size, bias, stride, padding, dilation);
       }
@@ -788,12 +764,6 @@ at::Tensor _convolution(
           bias,
           params.stride,
           params.padding);
-#ifdef USE_VULKAN
-  } else if (params.use_vulkan(input, weight)) {
-    output = at::native::vulkan_convolution(
-        input, weight, bias,
-        params.padding, params.stride, params.dilation, params.groups);
-#endif
   } else if (input.device().type() == c10::DeviceType::CPU || input.device().type() == c10::DeviceType::CUDA) {
     if (params.groups == 1) {
       output = at::_convolution_nogroup(

diff --git a/aten/src/ATen/native/Copy.cpp b/aten/src/ATen/native/Copy.cpp
@@ -6,13 +6,11 @@
 #include <ATen/native/TensorIterator.h>
 #include <ATen/native/quantized/Copy.h>
 #include <ATen/quantized/Quantizer.h>
+#include <ATen/vulkan/Context.h>
 #include <ATen/MemoryOverlap.h>
 #include <ATen/NamedTensorUtils.h>
 #include <torch/library.h>
 
-#ifdef USE_VULKAN
-#include <ATen/native/vulkan/VulkanAten.h>
-#endif
 namespace {
 
 using namespace at;
@@ -131,11 +129,9 @@ static Tensor & copy_impl(Tensor & self, const Tensor & src, bool non_blocking)
     TORCH_CHECK(false, "Copying from quantized Tensor to non-quantized Tensor is not allowed, please use dequantize to get a float Tensor from a quantized Tensor");
   }
 
-#ifdef USE_VULKAN
   if (self.device().type() == at::kVulkan || src.device().type() == at::kVulkan) {
-    return vulkan_copy_(self, src);
+    return at::vulkan::vulkan_copy_(self, src);
   }
-#endif
 
   auto iter = TensorIteratorConfig()
     .set_check_mem_overlap(true)

diff --git a/aten/src/ATen/native/Pooling.cpp b/aten/src/ATen/native/Pooling.cpp
@@ -8,9 +8,6 @@
 #include <c10/util/Exception.h>
 
 #include <tuple>
-#ifdef USE_VULKAN
-#include <ATen/native/vulkan/VulkanAten.h>
-#endif
 
 namespace at { namespace native {
 
@@ -138,13 +135,6 @@ Tensor max_pool2d(
         self, kernel_size, stride, padding, dilation, ceil_mode);
   }
 
-#ifdef USE_VULKAN
-  if (self.is_vulkan()) {
-    return at::native::vulkan_max_pool2d(
-      self, kernel_size, stride, padding, dilation, ceil_mode);
-  }
-#endif
-
 #if defined(C10_MOBILE)
   if(xnnpack::use_max_pool2d(self, kernel_size, padding, stride,
                              dilation, ceil_mode)) {

diff --git a/aten/src/ATen/native/TensorShape.cpp b/aten/src/ATen/native/TensorShape.cpp
@@ -18,10 +18,6 @@
 #include <ATen/native/Copy.h>
 #include <ATen/MemoryOverlap.h>
 
-#ifdef USE_VULKAN
-#include <ATen/native/vulkan/VulkanAten.h>
-#endif
-
 namespace at {
 namespace native {
 
@@ -843,11 +839,6 @@ Tensor reshape(const Tensor& self, IntArrayRef proposed_shape) {
   if (self.is_mkldnn()) {
     return at::_mkldnn_reshape(self, shape);
   }
-#ifdef USE_VULKAN
-    if (self.is_vulkan()) {
-      return at::native::vulkan_reshape(self, shape);
-    }
-#endif
 
   auto stride =
       at::detail::computeStride(self.sizes(), self.strides(), shape);

diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -303,7 +303,6 @@
     CPU, CUDA: add
     SparseCPU, SparseCUDA: add_sparse
     MkldnnCPU: mkldnn_add
-    Vulkan: vulkan_add
 
 - func: add_.Tensor(Tensor(a!) self, Tensor other, *, Scalar alpha=1) -> Tensor(a!)
   variants: method
@@ -775,7 +774,6 @@
   dispatch:
     CPU, CUDA: clamp
     QuantizedCPU: clamp_quantized_cpu
-    Vulkan: vulkan_clamp
 
 - func: clamp_(Tensor(a!) self, Scalar? min=None, Scalar? max=None) -> Tensor(a!)
   variants: function, method
@@ -1188,7 +1186,6 @@
     CUDA: empty_cuda
     MkldnnCPU: empty_mkldnn
     SparseCPU, SparseCUDA: empty_sparse
-    Vulkan: empty_vulkan
 
 - func: new_empty(Tensor self, int[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
   use_c10_dispatcher: full
@@ -1244,7 +1241,6 @@
   dispatch:
     CPU: empty_strided_cpu
     CUDA: empty_strided_cuda
-    Vulkan: empty_strided_vulkan
 
 - func: erf(Tensor self) -> Tensor
   use_c10_dispatcher: full
@@ -1926,7 +1922,6 @@
   dispatch:
     CPU, CUDA: mean_cpu_gpu
     QuantizedCPU: mean_quantized_cpu
-    Vulkan: mean_vulkan
 
 - func: mean.out(Tensor self, int[1] dim, bool keepdim=False, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!)
   dispatch:
@@ -2073,7 +2068,6 @@
     CPU: mm_cpu
     CUDA: mm_cuda
     SparseCPU, SparseCUDA: _sparse_mm
-    Vulkan: vulkan_mm
 
 - func: mm.out(Tensor self, Tensor mat2, *, Tensor(a!) out) -> Tensor(a!)
   dispatch:
@@ -3441,7 +3435,6 @@
     CUDA: addmm_cuda
     SparseCPU: addmm_sparse_dense_cpu
     SparseCUDA: addmm_sparse_dense_cuda
-    Vulkan: vulkan_addmm
 
 - func: addmm_(Tensor(a!) self, Tensor mat1, Tensor mat2, *, Scalar beta=1, Scalar alpha=1) -> Tensor(a!)
   variants: method
@@ -5772,7 +5765,6 @@
   dispatch:
     CPU, CUDA: hardtanh_
     QuantizedCPU: hardtanh_quantized_cpu_
-    Vulkan: vulkan_hardtanh_
 
 - func: hardswish.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)
   python_module: nn
@@ -6510,7 +6502,6 @@
     CPU: upsample_nearest2d_cpu
     CUDA: upsample_nearest2d_cuda
     QuantizedCPU: upsample_nearest2d_quantized_cpu
-    Vulkan: upsample_nearest2d_vulkan
 
 - func: upsample_nearest2d_backward.grad_input(Tensor grad_output, int[2] output_size, int[4] input_size, float? scales_h=None, float? scales_w=None, *, Tensor(a!) grad_input) -> Tensor(a!)
   python_module: nn