Part pipeline: Fix wave size issue in vertex shader

The default wave size for a part pipeline compilation rely on whether
sugroup size is used in the other part. In the whole pipeline mode it
is normally handled by querying the shader modes field, so we need to
make sure that the shader modes have proper information even in the part
pipeline mode.

lgc/builder/Builder.cpp

@@ -120,6 +120,15 @@ void Builder::setComputeShaderMode(const ComputeShaderMode &computeShaderMode) {
   getShaderModes()->setComputeShaderMode(computeShaderMode);
 }
 
+// =====================================================================================================================
+// Set subgroup size usage
+//
+// @param stage : Shader stage
+// @param usage : Subgroup size usage
+void Builder::setSubgroupSizeUsage(ShaderStage stage, bool usage){
+  getShaderModes()->setSubgroupSizeUsage(stage, usage);
+}
+
 // =====================================================================================================================
 // Get the compute shader mode (workgroup size)
 const ComputeShaderMode &Builder::getComputeShaderMode() {

lgc/include/lgc/state/ShaderModes.h

@@ -74,6 +74,9 @@ class ShaderModes {
   // Get the compute shader mode (workgroup size)
   const ComputeShaderMode &getComputeShaderMode();
 
+  // Set subgroup size usage
+  void setSubgroupSizeUsage(ShaderStage stage, bool usage);
+
   // Clear all modes
   void clear();
 

lgc/interface/lgc/Builder.h

@@ -184,6 +184,9 @@ class Builder : public BuilderCommon {
   // add more fields. A local struct variable can be zero-initialized with " = {}".
   void setComputeShaderMode(const ComputeShaderMode &computeShaderMode);
 
+  // Set subgroup size usage
+  void setSubgroupSizeUsage(ShaderStage stage, bool usage);
+
   // Get the compute shader mode (workgroup size)
   const ComputeShaderMode &getComputeShaderMode();
 

lgc/state/ShaderModes.cpp

@@ -159,6 +159,16 @@ const ComputeShaderMode &ShaderModes::getComputeShaderMode() {
   return m_computeShaderMode;
 }
 
+// =====================================================================================================================
+// Set subgroup size usage
+//
+// @param stage : Shader stage
+// @param usage : Subgroup size usage
+void ShaderModes::setSubgroupSizeUsage(ShaderStage stage, bool usage){
+  auto modes = MutableArrayRef<CommonShaderMode>(m_commonShaderModes);
+  modes[stage].useSubgroupSize = usage;
+}
+
 // =====================================================================================================================
 // Record shader modes (common and specific) into IR metadata
 //

llpc/context/llpcComputeContext.cpp

@@ -61,10 +61,10 @@ const PipelineShaderInfo *ComputeContext::getPipelineShaderInfo(ShaderStage shad
 }
 
 // =====================================================================================================================
-// Check whether the pipeline uses features relevant to subgroup size
-bool ComputeContext::usesSubgroupSize() const {
+// Gets subgroup size usage
+unsigned ComputeContext::getSubgroupSizeUsage() const {
   const ShaderModuleData *moduleData = reinterpret_cast<const ShaderModuleData *>(m_pipelineInfo->cs.pModuleData);
-  return moduleData->usage.useSubgroupSize;
+  return moduleData->usage.useSubgroupSize << ShaderStageCompute;
 }
 
 } // namespace Llpc

llpc/context/llpcComputeContext.h

@@ -62,8 +62,8 @@ class ComputeContext : public PipelineContext {
   // Gets the count of active shader stages
   virtual unsigned getActiveShaderStageCount() const { return 1; }
 
-  // Check whether the pipeline uses features relevant to subgroup size
-  virtual bool usesSubgroupSize() const;
+  // Gets subgroup size usage
+  virtual unsigned getSubgroupSizeUsage() const;
 
   // Gets per pipeline options
   virtual const PipelineOptions *getPipelineOptions() const { return &m_pipelineInfo->options; }

llpc/context/llpcGraphicsContext.cpp

@@ -118,8 +118,8 @@ const PipelineShaderInfo *GraphicsContext::getPipelineShaderInfo(ShaderStage sha
 }
 
 // =====================================================================================================================
-// Check whether the pipeline uses features relevant to subgroup size
-bool GraphicsContext::usesSubgroupSize() const {
+// Gets subgroup size usage
+unsigned GraphicsContext::getSubgroupSizeUsage() const {
   // clang-format off
   std::array<const PipelineShaderInfo *, ShaderStageGfxCount> shaderInfos = {
     &m_pipelineInfo->vs,
@@ -129,15 +129,16 @@ bool GraphicsContext::usesSubgroupSize() const {
     &m_pipelineInfo->fs,
   };
   // clang-format on
-  for (auto shaderInfo : shaderInfos) {
+  unsigned bitmask = 0;
+  for (unsigned shaderInfoIdx = 0; shaderInfoIdx < shaderInfos.size(); ++shaderInfoIdx) {
+    auto shaderInfo = shaderInfos[shaderInfoIdx];
     if (!shaderInfo->pModuleData)
       continue;
     auto *moduleData = reinterpret_cast<const ShaderModuleData *>(shaderInfo->pModuleData);
-    if (!moduleData->usage.useSubgroupSize)
-      continue;
-    return true;
+    if (moduleData->usage.useSubgroupSize)
+      bitmask |= (1 << shaderInfoIdx);
   }
-  return false;
+  return bitmask;
 }
 
 } // namespace Llpc

llpc/context/llpcGraphicsContext.h

@@ -68,8 +68,8 @@ class GraphicsContext : public PipelineContext {
   // Gets per pipeline options
   virtual const PipelineOptions *getPipelineOptions() const { return &m_pipelineInfo->options; }
 
-  // Check whether the pipeline uses features relevant to subgroup size
-  virtual bool usesSubgroupSize() const;
+  // Gets subgroup size usage
+  virtual unsigned getSubgroupSizeUsage() const;
 
 private:
   GraphicsContext() = delete;

llpc/context/llpcPipelineContext.h

@@ -103,8 +103,9 @@ class PipelineContext {
 
   // Checks whether the pipeline is graphics or compute
   virtual bool isGraphics() const = 0;
-  // Check whether the pipeline uses features relevant to subgroup size
-  virtual bool usesSubgroupSize() const = 0;
+
+  // Gets subgroup size usage - bitmask with one bit per stage.
+  virtual unsigned getSubgroupSizeUsage() const = 0;
 
   // Gets pipeline shader info of the specified shader stage
   virtual const PipelineShaderInfo *getPipelineShaderInfo(ShaderStage shaderStage) const = 0;

llpc/translator/lib/SPIRV/SPIRVReader.cpp

@@ -6514,6 +6514,16 @@ bool SPIRVToLLVM::translate(ExecutionModel entryExecModel, const char *entryName
 
   shaderMode.useSubgroupSize = m_moduleUsage->useSubgroupSize;
 
+  // Shader modes contain also data for other modules (subgroup size usage), so query it in the pipeline context.
+  auto pipelineContext = (static_cast<Llpc::Context *>(m_context))->getPipelineContext();
+  unsigned subgroupSizeUsage = pipelineContext->getSubgroupSizeUsage();
+  if (pipelineContext->isGraphics() && subgroupSizeUsage){
+    for(lgc::ShaderStage stage : lgc::enumRange<lgc::ShaderStage>()) {
+      if (subgroupSizeUsage & (1 << stage)){
+        getBuilder()->setSubgroupSizeUsage(stage, true);
+      }
+    }
+  }
   getBuilder()->setCommonShaderMode(shaderMode);
 
   m_enableXfb = m_bm->getCapability().find(CapabilityTransformFeedback) != m_bm->getCapability().end();