[vulkan] Use automatically generated descriptor set layouts (pytorch#…

…81716) Differential Revision: [D37966837](https://our.internmc.facebook.com/intern/diff/D37966837/) Pull Request resolved: pytorch#81716 Approved by: https://github.com/kirklandsign
bmedishe · Jul 20, 2022 · e85bdd5 · e85bdd5
1 parent 96958be
commit e85bdd5
Show file tree

Hide file tree

Showing 25 changed files with 40 additions and 257 deletions.
diff --git a/aten/src/ATen/native/vulkan/api/Common.h b/aten/src/ATen/native/vulkan/api/Common.h
@@ -14,9 +14,9 @@
   }
 #else
 #include <ATen/native/vulkan/spv.h>
-#define VK_KERNEL(name)                     \
-  ::at::native::vulkan::api::ShaderSource { \
-#name, name##_spv, name##_spv_len,      \
+#define VK_KERNEL(name)                                  \
+  ::at::native::vulkan::api::ShaderSource {              \
+#name, name##_spv, name##_spv_len, name##_spv_layout \
   }
 #endif /* USE_VULKAN_SHADERC_RUNTIME */
 

diff --git a/aten/src/ATen/native/vulkan/api/Context.cpp b/aten/src/ATen/native/vulkan/api/Context.cpp
@@ -42,11 +42,10 @@ Context::~Context() {
 
 DescriptorSet Context::submit_compute_prologue(
     CommandBuffer& command_buffer,
-    const ShaderLayout::Signature& shader_layout_signature,
     const ShaderSource& shader_descriptor,
     const utils::uvec3& local_workgroup_size) {
   const VkDescriptorSetLayout shader_layout =
-      shader_layout_cache().retrieve(shader_layout_signature);
+      shader_layout_cache().retrieve(shader_descriptor.kernel_layout);
 
   const VkPipelineLayout pipeline_layout =
       pipeline_layout_cache().retrieve(shader_layout);
@@ -59,7 +58,7 @@ DescriptorSet Context::submit_compute_prologue(
   command_buffer.bind_pipeline(pipeline, pipeline_layout, local_workgroup_size);
 
   return descriptor_pool().get_descriptor_set(
-      shader_layout, shader_layout_signature);
+      shader_layout, shader_descriptor.kernel_layout);
 }
 
 void Context::submit_compute_epilogue(

diff --git a/aten/src/ATen/native/vulkan/api/Context.h b/aten/src/ATen/native/vulkan/api/Context.h
@@ -153,7 +153,6 @@ class Context final {
 
   DescriptorSet submit_compute_prologue(
       CommandBuffer&,
-      const ShaderLayout::Signature&,
       const ShaderSource&,
       const utils::uvec3&);
 
@@ -166,7 +165,6 @@ class Context final {
  public:
   template <typename... Arguments>
   void submit_compute_job(
-      const ShaderLayout::Signature&,
       const ShaderSource&,
       const PipelineBarrier&,
       const utils::uvec3&,
@@ -270,7 +268,6 @@ inline void bind(
 
 template <typename... Arguments>
 inline void Context::submit_compute_job(
-    const ShaderLayout::Signature& shader_layout_signature,
     const ShaderSource& shader_descriptor,
     const PipelineBarrier& pipeline_barrier,
     const utils::uvec3& global_work_group,
@@ -302,8 +299,8 @@ inline void Context::submit_compute_job(
 #endif /* USE_VULKAN_GPU_DIAGNOSTICS */
 
   // Factor out template parameter independent code to minimize code bloat.
-  DescriptorSet descriptor_set = submit_compute_prologue(
-      cmd_, shader_layout_signature, shader_descriptor, local_work_group_size);
+  DescriptorSet descriptor_set =
+      submit_compute_prologue(cmd_, shader_descriptor, local_work_group_size);
 
   detail::bind(
       descriptor_set,

diff --git a/aten/src/ATen/native/vulkan/api/Shader.cpp b/aten/src/ATen/native/vulkan/api/Shader.cpp
@@ -27,7 +27,8 @@ ShaderSource::ShaderSource(std::string name, const char* const glsl_src)
 ShaderSource::ShaderSource(
     std::string name,
     const uint32_t* const spirv_bin,
-    const uint32_t size)
+    const uint32_t size,
+    const std::vector<VkDescriptorType>& layout)
     : type(Type::SPIRV),
       src_code{
           .spirv =
@@ -36,7 +37,8 @@ ShaderSource::ShaderSource(
                   size,
               },
       },
-      kernel_name{std::move(name)} {}
+      kernel_name{std::move(name)},
+      kernel_layout{layout} {}
 
 bool operator==(const ShaderSource& _1, const ShaderSource& _2) {
   if (_1.type != _2.type) {

diff --git a/aten/src/ATen/native/vulkan/api/Shader.h b/aten/src/ATen/native/vulkan/api/Shader.h
@@ -11,28 +11,6 @@ namespace native {
 namespace vulkan {
 namespace api {
 
-struct ShaderSource final {
-  enum class Type { GLSL, SPIRV } type;
-
-  union {
-    struct {
-      const char* src; // Null-terminated
-      uint32_t unused; // padding
-    } glsl;
-    struct {
-      const uint32_t* bin;
-      uint32_t size;
-    } spirv;
-  } src_code;
-
-  std::string kernel_name;
-  explicit ShaderSource(std::string name, const char* glsl);
-  explicit ShaderSource(
-      std::string name,
-      const uint32_t* spirv,
-      uint32_t bytes);
-};
-
 class ShaderLayout final {
  public:
   using Signature = c10::SmallVector<VkDescriptorType, 6u>;
@@ -62,6 +40,31 @@ class ShaderLayout final {
   friend void swap(ShaderLayout& lhs, ShaderLayout& rhs) noexcept;
 };
 
+struct ShaderSource final {
+  enum class Type { GLSL, SPIRV } type;
+
+  union {
+    struct {
+      const char* src; // Null-terminated
+      uint32_t unused; // padding
+    } glsl;
+    struct {
+      const uint32_t* bin;
+      uint32_t size;
+    } spirv;
+  } src_code;
+
+  std::string kernel_name;
+  ShaderLayout::Signature kernel_layout;
+
+  explicit ShaderSource(std::string, const char*);
+  explicit ShaderSource(
+      std::string,
+      const uint32_t*,
+      const uint32_t,
+      const std::vector<VkDescriptorType>&);
+};
+
 class ShaderModule final {
  public:
   explicit ShaderModule(const VkDevice device, const ShaderSource& source);

diff --git a/aten/src/ATen/native/vulkan/ops/Arithmetic.cpp b/aten/src/ATen/native/vulkan/ops/Arithmetic.cpp
@@ -80,12 +80,6 @@ Tensor arithmetic_scalar(
   api::PipelineBarrier pipeline_barrier{};
 
   context->submit_compute_job(
-      // shader layout signature
-      {
-          VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
-          VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
-          VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
-      },
       // shader descriptor
       shader_descriptor,
       // pipeline barrier
@@ -135,11 +129,6 @@ Tensor& arithmetic_scalar_(
   api::PipelineBarrier pipeline_barrier{};
 
   context->submit_compute_job(
-      // shader layout signature
-      {
-          VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
-          VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
-      },
       // shader descriptor
       shader_descriptor,
       // pipeline barrier
@@ -202,13 +191,6 @@ Tensor arithmetic_tensor(
   api::PipelineBarrier pipeline_barrier{};
 
   context->submit_compute_job(
-      // shader layout signature
-      {
-          VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
-          VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
-          VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
-          VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
-      },
       // shader descriptor
       shader_descriptor,
       // pipeline barrier
@@ -267,12 +249,6 @@ Tensor& arithmetic_tensor_(
   api::PipelineBarrier pipeline_barrier{};
 
   context->submit_compute_job(
-      // shader layout signature
-      {
-          VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
-          VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
-          VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
-      },
       // shader descriptor
       shader_descriptor,
       // pipeline barrier

diff --git a/aten/src/ATen/native/vulkan/ops/Batchnorm.cpp b/aten/src/ATen/native/vulkan/ops/Batchnorm.cpp
@@ -106,16 +106,6 @@ Tensor batch_norm(
   api::PipelineBarrier pipeline_barrier{};
 
   context->submit_compute_job(
-      // shader layout signature
-      {
-          VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
-          VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
-          VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
-          VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
-          VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
-          VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
-          VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
-      },
       // shader descriptor
       VK_KERNEL(batchnorm),
       // pipeline barrier

diff --git a/aten/src/ATen/native/vulkan/ops/Clamp.cpp b/aten/src/ATen/native/vulkan/ops/Clamp.cpp
@@ -44,12 +44,6 @@ Tensor _clamp(
   api::PipelineBarrier pipeline_barrier{};
 
   context->submit_compute_job(
-      // shader layout signature
-      {
-          VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
-          VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
-          VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
-      },
       // shader descriptor
       shader_descriptor,
       // pipeline barrier
@@ -112,11 +106,6 @@ Tensor& _clamp_(
   api::PipelineBarrier pipeline_barrier{};
 
   context->submit_compute_job(
-      // shader layout signature
-      {
-          VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
-          VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
-      },
       // shader descriptor
       shader_descriptor,
       // pipeline barrier
@@ -178,12 +167,6 @@ Tensor activation(
   api::PipelineBarrier pipeline_barrier{};
 
   context->submit_compute_job(
-      // shader layout signature
-      {
-          VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
-          VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
-          VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
-      },
       // shader descriptor
       shader_descriptor,
       // pipeline barrier
@@ -229,11 +212,6 @@ Tensor& activation_(
   api::PipelineBarrier pipeline_barrier{};
 
   context->submit_compute_job(
-      // shader layout signature
-      {
-          VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
-          VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
-      },
       // shader descriptor
       shader_descriptor,
       // pipeline barrier
@@ -316,12 +294,6 @@ Tensor activation_scalar(
   api::PipelineBarrier pipeline_barrier{};
 
   context->submit_compute_job(
-      // shader layout signature
-      {
-          VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
-          VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
-          VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
-      },
       // shader descriptor
       shader_descriptor,
       // pipeline barrier
@@ -370,11 +342,6 @@ Tensor& activation_scalar_(
   api::PipelineBarrier pipeline_barrier{};
 
   context->submit_compute_job(
-      // shader layout signature
-      {
-          VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
-          VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
-      },
       // shader descriptor
       shader_descriptor,
       // pipeline barrier

diff --git a/aten/src/ATen/native/vulkan/ops/Concat.cpp b/aten/src/ATen/native/vulkan/ops/Concat.cpp
@@ -61,12 +61,6 @@ Tensor cat_feature(const TensorList tensors, vTensor& v_output) {
     api::PipelineBarrier pipeline_barrier{};
 
     context->submit_compute_job(
-        // shader layout signature
-        {
-            VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
-            VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
-            VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
-        },
         // shader descriptor
         VK_KERNEL(cat_feature),
         // pipeline barrier

diff --git a/aten/src/ATen/native/vulkan/ops/Convolution.cpp b/aten/src/ATen/native/vulkan/ops/Convolution.cpp
@@ -481,14 +481,6 @@ void conv2d_sliding_window(
   api::PipelineBarrier pipeline_barrier{};
 
   context->submit_compute_job(
-      // shader layout signature
-      {
-          VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
-          VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
-          VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
-          VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
-          VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
-      },
       // shader descriptor
       shader,
       // pipeline barrier

diff --git a/aten/src/ATen/native/vulkan/ops/Glu.cpp b/aten/src/ATen/native/vulkan/ops/Glu.cpp
@@ -45,12 +45,6 @@ Tensor glu(const at::Tensor& input_arg, const int64_t dim = -1) {
   api::PipelineBarrier pipeline_barrier{};
 
   context->submit_compute_job(
-      // shader layout signature
-      {
-          VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
-          VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
-          VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
-      },
       // shader descriptor
       VK_KERNEL(glu),
       // pipeline barrier

diff --git a/aten/src/ATen/native/vulkan/ops/Layernorm.cpp b/aten/src/ATen/native/vulkan/ops/Layernorm.cpp
@@ -114,14 +114,6 @@ Tensor layer_norm(
   api::PipelineBarrier pipeline_barrier{};
 
   context->submit_compute_job(
-      // shader layout signature
-      {
-          VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
-          VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
-          VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
-          VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
-          VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
-      },
       // shader descriptor
       VK_KERNEL(layernorm),
       // pipeline barrier