Tweak Vulkan memory use. (#47728)

Summary: Pull Request resolved: #47728

Test Plan: Imported from OSS

Reviewed By: SS-JIA

Differential Revision: D25032740

Pulled By: AshkanAliabadi

fbshipit-source-id: 7eb72538dc1aa3feb4e2f8c4ff9c675eb8e97057

CMakeLists.txt

@@ -205,10 +205,10 @@ cmake_dependent_option(
     USE_VALGRIND "Use Valgrind. Only available on Linux." ON
     "LINUX" OFF)
 option(USE_VULKAN "Use Vulkan GPU backend" OFF)
-option(USE_VULKAN_API "Use Vulkan GPU backend v2" OFF)
-option(USE_VULKAN_WRAPPER "Use Vulkan wrapper" ON)
-option(USE_VULKAN_SHADERC_RUNTIME "Use Vulkan Shader compilation runtime(Needs shaderc lib)" OFF)
-option(USE_VULKAN_RELAXED_PRECISION "Use Vulkan relaxed precision(mediump)" OFF)
+option(USE_VULKAN_FP16_INFERENCE "Vulkan - Use fp16 inference even on fp32 tensors" ON)
+option(USE_VULKAN_RELAXED_PRECISION "Vulkan - Use relaxed precision math in the kernels (mediump)" OFF)
+option(USE_VULKAN_SHADERC_RUNTIME "Vulkan - Use runtime shader compilation (needs libshaderc)" OFF)
+option(USE_VULKAN_WRAPPER "Vulkan - Dynamically load Vulkan functions" ON)
 option(USE_XNNPACK "Use XNNPACK" ON)
 option(USE_ZMQ "Use ZMQ" OFF)
 option(USE_ZSTD "Use ZSTD" OFF)
@@ -554,22 +554,23 @@ endif()
 
 if(USE_VULKAN)
   string(APPEND CMAKE_CXX_FLAGS " -DUSE_VULKAN")
-endif()
-
-if(USE_VULKAN_API)
   string(APPEND CMAKE_CXX_FLAGS " -DUSE_VULKAN_API")
-endif()
 
-if(USE_VULKAN_WRAPPER)
-  string(APPEND CMAKE_CXX_FLAGS " -DUSE_VULKAN_WRAPPER")
-endif()
+  if(USE_VULKAN_FP16_INFERENCE)
+    string(APPEND CMAKE_CXX_FLAGS " -DUSE_VULKAN_FP16_INFERENCE")
+  endif()
 
-if(USE_VULKAN_SHADERC_RUNTIME)
-  string(APPEND CMAKE_CXX_FLAGS " -DUSE_VULKAN_SHADERC_RUNTIME")
-endif()
+  if(USE_VULKAN_RELAXED_PRECISION)
+    string(APPEND CMAKE_CXX_FLAGS " -DUSE_VULKAN_RELAXED_PRECISION")
+  endif()
 
-if(USE_VULKAN_RELAXED_PRECISION)
-  string(APPEND CMAKE_CXX_FLAGS " -DUSE_VULKAN_RELAXED_PRECISION")
+  if(USE_VULKAN_SHADERC_RUNTIME)
+    string(APPEND CMAKE_CXX_FLAGS " -DUSE_VULKAN_SHADERC_RUNTIME")
+  endif()
+
+  if(USE_VULKAN_WRAPPER)
+    string(APPEND CMAKE_CXX_FLAGS " -DUSE_VULKAN_WRAPPER")
+  endif()
 endif()
 
 if(USE_PYTORCH_METAL)

aten/src/ATen/CMakeLists.txt

@@ -65,8 +65,7 @@ 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_api_cpp "native/vulkan/api/*.cpp" "native/vulkan/ops/*.cpp")
+file(GLOB native_vulkan_cpp "native/vulkan/*.cpp" "native/vulkan/api/*.cpp" "native/vulkan/ops/*.cpp")
 
 # Metal
 file(GLOB metal_h "metal/*.h")
@@ -126,9 +125,6 @@ if(AT_MKLDNN_ENABLED)
 endif()
 if(USE_VULKAN)
   set(all_cpu_cpp ${all_cpu_cpp} ${vulkan_cpp} ${native_vulkan_cpp} ${vulkan_generated_cpp})
-  if(USE_VULKAN_API)
-    set(all_cpu_cpp ${all_cpu_cpp} ${native_vulkan_api_cpp})
-  endif()
 else()
   set(all_cpu_cpp ${all_cpu_cpp} ${vulkan_cpp})
 endif()

aten/src/ATen/native/vulkan/Vulkan.cpp

@@ -779,17 +779,17 @@ void ComputeUnit::createComputePipeline(
   {
     uint32_t offset = 0;
     size_t size = sizeof(WorkGroupSize::x);
-    spMapEntries[0].constantID = 1;
+    spMapEntries[0].constantID = 0;
     spMapEntries[0].offset = offset;
     spMapEntries[0].size = size;
     offset += size;
     size = sizeof(WorkGroupSize::y);
-    spMapEntries[1].constantID = 2;
+    spMapEntries[1].constantID = 1;
     spMapEntries[1].offset = offset;
     spMapEntries[1].size = size;
     offset += size;
     size = sizeof(WorkGroupSize::z);
-    spMapEntries[2].constantID = 3;
+    spMapEntries[2].constantID = 2;
     spMapEntries[2].offset = offset;
     spMapEntries[2].size = size;
   }

aten/src/ATen/native/vulkan/api/Adapter.h

@@ -33,7 +33,7 @@ struct Adapter final {
   }
 
   inline Shader::WorkGroup local_work_group_size() const {
-    return { 8u, 8u, 1u, };
+    return { 4u, 4u, 4u, };
   }
 };
 

aten/src/ATen/native/vulkan/api/Allocator.h

@@ -13,6 +13,8 @@
 
 #ifdef USE_VULKAN_WRAPPER
   #define VMA_STATIC_VULKAN_FUNCTIONS 0
+#else
+  #define VMA_DYNAMIC_VULKAN_FUNCTIONS 0
 #endif
 
 #ifdef DEBUG

aten/src/ATen/native/vulkan/api/Cache.h

@@ -72,7 +72,7 @@ template<typename Factory>
 inline auto Cache<Factory>::retrieve(
     const Descriptor& descriptor) {
   auto iterator = cache_.find(descriptor);
-  if (cache_.cend() == iterator) {
+  if C10_UNLIKELY(cache_.cend() == iterator) {
     iterator = cache_.insert({descriptor, factory_(descriptor)}).first;
   }