igl | vulkan | Check if SPIR-V shader uses push constants

Summary: Added a new member field `SpvModuleInfo::hasPushConstants` to check if a SPIR-V shader uses push constants.

Reviewed By: mmaurer

Differential Revision: D52446637

fbshipit-source-id: 4b05638c04564cc56f9469176ca8f4d308807d32

src/igl/vulkan/PipelineState.cpp

@@ -15,14 +15,21 @@
 
 namespace igl::vulkan {
 
-void PipelineState::initializeSpvModuleInfoFromShaderStages(IShaderStages* stages) {
+void PipelineState::initializeSpvModuleInfoFromShaderStages(const VulkanContext& ctx,
+                                                            IShaderStages* stages) {
   auto* smComp = static_cast<igl::vulkan::ShaderModule*>(stages->getComputeModule().get());
 
+  VkShaderStageFlags pushConstantMask = 0;
+
   if (smComp) {
     // compute
     ensureShaderModule(smComp);
 
     info_ = smComp->getVulkanShaderModule().getSpvModuleInfo();
+
+    if (info_.hasPushConstants) {
+      pushConstantMask |= VK_SHADER_STAGE_COMPUTE_BIT;
+    }
   } else {
     auto* smVert = static_cast<igl::vulkan::ShaderModule*>(stages->getVertexModule().get());
     auto* smFrag = static_cast<igl::vulkan::ShaderModule*>(stages->getFragmentModule().get());
@@ -34,16 +41,37 @@ void PipelineState::initializeSpvModuleInfoFromShaderStages(IShaderStages* stage
     const util::SpvModuleInfo& infoVert = smVert->getVulkanShaderModule().getSpvModuleInfo();
     const util::SpvModuleInfo& infoFrag = smFrag->getVulkanShaderModule().getSpvModuleInfo();
 
+    if (infoVert.hasPushConstants) {
+      pushConstantMask |= VK_SHADER_STAGE_VERTEX_BIT;
+    }
+    if (infoFrag.hasPushConstants) {
+      pushConstantMask |= VK_SHADER_STAGE_FRAGMENT_BIT;
+    }
+
     info_ = util::mergeReflectionData(infoVert, infoFrag);
   }
+
+  if (pushConstantMask) {
+    const VkPhysicalDeviceLimits& limits = ctx.getVkPhysicalDeviceProperties().limits;
+
+    constexpr uint32_t kPushConstantsSize = 128;
+
+    if (!IGL_VERIFY(kPushConstantsSize <= limits.maxPushConstantsSize)) {
+      IGL_LOG_ERROR("Push constants size exceeded %u (max %u bytes)",
+                    kPushConstantsSize,
+                    limits.maxPushConstantsSize);
+    }
+
+    pushConstantRange_ = ivkGetPushConstantRange(pushConstantMask, 0, kPushConstantsSize);
+  }
 }
 
 PipelineState::PipelineState(const VulkanContext& ctx,
                              IShaderStages* stages,
                              const char* debugName) {
   IGL_ASSERT(stages);
 
-  initializeSpvModuleInfoFromShaderStages(stages);
+  initializeSpvModuleInfoFromShaderStages(ctx, stages);
 
   // Create all Vulkan descriptor set layouts for this pipeline
 

src/igl/vulkan/PipelineState.h

@@ -33,13 +33,13 @@ class PipelineState {
   }
 
  private:
-  void initializeSpvModuleInfoFromShaderStages(IShaderStages* stages);
-
- protected:
-  friend class ResourcesBinder;
+  void initializeSpvModuleInfoFromShaderStages(const VulkanContext& ctx, IShaderStages* stages);
 
+ public:
   igl::vulkan::util::SpvModuleInfo info_;
 
+  VkPushConstantRange pushConstantRange_ = {};
+
   mutable std::unique_ptr<igl::vulkan::VulkanPipelineLayout> pipelineLayout_;
 
   // the last seen VkDescriptorSetLayout from VulkanContext::dslBindless_

src/igl/vulkan/VulkanHelpers.c

@@ -1170,7 +1170,7 @@ VkPipelineLayoutCreateInfo ivkGetPipelineLayoutCreateInfo(uint32_t numLayouts,
       .sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
       .setLayoutCount = numLayouts,
       .pSetLayouts = layouts,
-      .pushConstantRangeCount = 1,
+      .pushConstantRangeCount = range ? 1u : 0u,
       .pPushConstantRanges = range,
   };
   return ci;

src/igl/vulkan/util/SpvReflection.cpp

@@ -228,6 +228,9 @@ SpvModuleInfo getReflectionData(const uint32_t* spirv, size_t numBytes) {
         break;
       }
     }
+    if (id.opCode == SpvOpVariable && id.storageClass == SpvStorageClassPushConstant) {
+      info.hasPushConstants = true;
+    }
   }
 
   return info;
@@ -260,6 +263,8 @@ SpvModuleInfo mergeReflectionData(const SpvModuleInfo& info1, const SpvModuleInf
   combineDescriptions(result.storageBuffers, info1.storageBuffers, info2.storageBuffers);
   combineDescriptions(result.textures, info1.textures, info2.textures);
 
+  result.hasPushConstants = info1.hasPushConstants || info2.hasPushConstants;
+
   return result;
 }
 

src/igl/vulkan/util/SpvReflection.h

@@ -32,6 +32,7 @@ struct SpvModuleInfo {
   std::vector<BufferDescription> uniformBuffers;
   std::vector<BufferDescription> storageBuffers;
   std::vector<TextureDescription> textures;
+  bool hasPushConstants = false;
 };
 
 SpvModuleInfo getReflectionData(const uint32_t* spirv, size_t numBytes);