Implementation of GroupSharedLimit to allow increased GroupSharedMemory (#7871)

Add a new HLSL attribute for Compute, Amp and Mesh shaders:
GroupSharedLimit.

This is used to limit the amount of group shared memory a shader is
allowed to statically declare, and validation will fail if the limit is
exceeded.

There is no upper limit on this attribute, and it is expected that
shader writers set the limit as the lowest common denominator for their
target hardware and software use case (typically 48k or 64k for modern
GPUs).

If no attribute is declared the existing 32k limit is used to be
compatible with existing shaders.

Extends the PSV structures to include the selected limit so that runtime
validation can reject the shader if it exceeds the device support.

---------

Co-authored-by: Jack Elliott <[email protected]>
Co-authored-by: Tex Riddell <[email protected]>

Author: 3 people

include/dxc/DXIL/DxilFunctionProps.h

@@ -117,6 +117,7 @@ struct DxilFunctionProps {
     memset(&Node, 0, sizeof(Node));
     Node.LaunchType = DXIL::NodeLaunchType::Invalid;
     Node.LocalRootArgumentsTableIndex = -1;
+    groupSharedLimitBytes = 0;
   }
   union {
     // Geometry shader.
@@ -174,6 +175,8 @@ struct DxilFunctionProps {
   // numThreads shared between multiple shader types and node shaders.
   unsigned numThreads[3];
 
+  unsigned groupSharedLimitBytes;
+
   struct NodeProps {
     DXIL::NodeLaunchType LaunchType = DXIL::NodeLaunchType::Invalid;
     bool IsProgramEntry;

include/dxc/DXIL/DxilMetadataHelper.h

@@ -320,6 +320,7 @@ class DxilMDHelper {
   static const unsigned kDxilNodeOutputsTag = 21;
   static const unsigned kDxilNodeMaxDispatchGridTag = 22;
   static const unsigned kDxilRangedWaveSizeTag = 23;
+  static const unsigned kDxilGroupSharedLimitTag = 24;
 
   // Node Input/Output State.
   static const unsigned kDxilNodeOutputIDTag = 0;

include/dxc/DXIL/DxilModule.h

@@ -254,6 +254,10 @@ class DxilModule {
   void SetNumThreads(unsigned x, unsigned y, unsigned z);
   unsigned GetNumThreads(unsigned idx) const;
 
+  unsigned GetGroupSharedLimit() const;
+  // The total amount of group shared memory (in bytes) used by the shader.
+  unsigned GetTGSMSizeInBytes() const;
+
   // Compute shader
   DxilWaveSize &GetWaveSize();
   const DxilWaveSize &GetWaveSize() const;

include/dxc/DxilContainer/DxilPipelineStateValidation.h

@@ -175,6 +175,10 @@ struct PSVRuntimeInfo3 : public PSVRuntimeInfo2 {
   uint32_t EntryFunctionName;
 };
 
+struct PSVRuntimeInfo4 : public PSVRuntimeInfo3 {
+  uint32_t NumBytesGroupSharedMemory;
+};
+
 enum class PSVResourceType {
   Invalid = 0,
 
@@ -474,7 +478,7 @@ class PSVSignatureElement {
              const uint32_t *SemanticIndexes) const;
 };
 
-#define MAX_PSV_VERSION 3
+#define MAX_PSV_VERSION 4
 
 struct PSVInitInfo {
   PSVInitInfo(uint32_t psvVersion) : PSVVersion(psvVersion) {}
@@ -491,7 +495,7 @@ struct PSVInitInfo {
   uint8_t SigPatchConstOrPrimVectors = 0;
   uint8_t SigOutputVectors[PSV_GS_MAX_STREAMS] = {0, 0, 0, 0};
 
-  static_assert(MAX_PSV_VERSION == 3, "otherwise this needs updating.");
+  static_assert(MAX_PSV_VERSION == 4, "otherwise this needs updating.");
   uint32_t RuntimeInfoSize() const {
     switch (PSVVersion) {
     case 0:
@@ -500,10 +504,12 @@ struct PSVInitInfo {
       return sizeof(PSVRuntimeInfo1);
     case 2:
       return sizeof(PSVRuntimeInfo2);
+    case 3:
+      return sizeof(PSVRuntimeInfo3);
     default:
       break;
     }
-    return sizeof(PSVRuntimeInfo3);
+    return sizeof(PSVRuntimeInfo4);
   }
   uint32_t ResourceBindInfoSize() const {
     if (PSVVersion < 2)
@@ -519,6 +525,7 @@ class DxilPipelineStateValidation {
   PSVRuntimeInfo1 *m_pPSVRuntimeInfo1 = nullptr;
   PSVRuntimeInfo2 *m_pPSVRuntimeInfo2 = nullptr;
   PSVRuntimeInfo3 *m_pPSVRuntimeInfo3 = nullptr;
+  PSVRuntimeInfo4 *m_pPSVRuntimeInfo4 = nullptr;
   uint32_t m_uResourceCount = 0;
   uint32_t m_uPSVResourceBindInfoSize = 0;
   void *m_pPSVResourceBindInfo = nullptr;
@@ -634,6 +641,8 @@ class DxilPipelineStateValidation {
 
   PSVRuntimeInfo3 *GetPSVRuntimeInfo3() const { return m_pPSVRuntimeInfo3; }
 
+  PSVRuntimeInfo4 *GetPSVRuntimeInfo4() const { return m_pPSVRuntimeInfo4; }
+
   uint32_t GetBindCount() const { return m_uResourceCount; }
 
   template <typename _T>
@@ -949,6 +958,8 @@ DxilPipelineStateValidation::ReadOrWrite(const void *pBits, uint32_t *pSize,
                 m_uPSVRuntimeInfoSize); // failure ok
   AssignDerived(&m_pPSVRuntimeInfo3, m_pPSVRuntimeInfo0,
                 m_uPSVRuntimeInfoSize); // failure ok
+  AssignDerived(&m_pPSVRuntimeInfo4, m_pPSVRuntimeInfo0,
+                m_uPSVRuntimeInfoSize); // failure ok
 
   // In RWMode::CalcSize, use temp runtime info to hold needed values from
   // initInfo
@@ -1137,11 +1148,13 @@ void SetupPSVInitInfo(PSVInitInfo &InitInfo, const DxilModule &DM);
 void SetShaderProps(PSVRuntimeInfo0 *pInfo, const DxilModule &DM);
 void SetShaderProps(PSVRuntimeInfo1 *pInfo1, const DxilModule &DM);
 void SetShaderProps(PSVRuntimeInfo2 *pInfo2, const DxilModule &DM);
+void SetShaderProps(PSVRuntimeInfo4 *pInfo4, const DxilModule &DM);
 
 void PrintPSVRuntimeInfo(llvm::raw_ostream &OS, PSVRuntimeInfo0 *pInfo0,
                          PSVRuntimeInfo1 *pInfo1, PSVRuntimeInfo2 *pInfo2,
-                         PSVRuntimeInfo3 *pInfo3, uint8_t ShaderKind,
-                         const char *EntryName, const char *Comment);
+                         PSVRuntimeInfo3 *pInfo3, PSVRuntimeInfo4 *pInfo4,
+                         uint8_t ShaderKind, const char *EntryName,
+                         const char *Comment);
 
 } // namespace hlsl
 

lib/DXIL/DxilMetadataHelper.cpp

@@ -1624,6 +1624,13 @@ MDTuple *DxilMDHelper::EmitDxilEntryProperties(uint64_t rawShaderFlag,
       }
       MDVals.emplace_back(MDNode::get(m_Ctx, WaveSizeVal));
     }
+
+    const hlsl::ShaderModel *SM = GetShaderModel();
+    if (SM->IsSMAtLeast(6, 10)) {
+      MDVals.emplace_back(
+          Uint32ToConstMD(DxilMDHelper::kDxilGroupSharedLimitTag));
+      MDVals.emplace_back(Uint32ToConstMD(props.groupSharedLimitBytes));
+    }
   } break;
     // Geometry shader.
   case DXIL::ShaderKind::Geometry: {
@@ -1773,6 +1780,13 @@ void DxilMDHelper::LoadDxilEntryProperties(const MDOperand &MDO,
       props.numThreads[2] = ConstMDToUint32(pNode->getOperand(2));
     } break;
 
+    case DxilMDHelper::kDxilGroupSharedLimitTag: {
+      DXASSERT(props.IsCS(), "else invalid shader kind");
+      props.groupSharedLimitBytes = ConstMDToUint32(MDO);
+      if (!m_pSM->IsSMAtLeast(6, 10))
+        m_bExtraMetadata = true;
+    } break;
+
     case DxilMDHelper::kDxilGSStateTag: {
       DXASSERT(props.IsGS(), "else invalid shader kind");
       auto &GS = props.ShaderProps.GS;

lib/DXIL/DxilModule.cpp

@@ -412,6 +412,28 @@ unsigned DxilModule::GetNumThreads(unsigned idx) const {
   return props.numThreads[idx];
 }
 
+unsigned DxilModule::GetGroupSharedLimit() const {
+  DXASSERT(m_DxilEntryPropsMap.size() == 1 &&
+               (m_pSM->IsCS() || m_pSM->IsMS() || m_pSM->IsAS()),
+           "only works for CS/MS/AS profiles");
+  const DxilFunctionProps &props = m_DxilEntryPropsMap.begin()->second->props;
+  DXASSERT_NOMSG(m_pSM->GetKind() == props.shaderKind);
+  return props.groupSharedLimitBytes;
+}
+
+unsigned DxilModule::GetTGSMSizeInBytes() const {
+  const DataLayout &DL = m_pModule->getDataLayout();
+  unsigned TGSMSize = 0;
+
+  for (GlobalVariable &GV : m_pModule->globals()) {
+    if (GV.getType()->getAddressSpace() == DXIL::kTGSMAddrSpace) {
+      TGSMSize += DL.getTypeAllocSize(GV.getType()->getElementType());
+    }
+  }
+
+  return TGSMSize;
+}
+
 DxilWaveSize &DxilModule::GetWaveSize() {
   return const_cast<DxilWaveSize &>(
       static_cast<const DxilModule *>(this)->GetWaveSize());

lib/DxilContainer/DxilContainerAssembler.cpp

@@ -798,6 +798,8 @@ class DxilPSVWriter : public DxilPartWriter {
     PSVRuntimeInfo1 *pInfo1 = m_PSV.GetPSVRuntimeInfo1();
     PSVRuntimeInfo2 *pInfo2 = m_PSV.GetPSVRuntimeInfo2();
     PSVRuntimeInfo3 *pInfo3 = m_PSV.GetPSVRuntimeInfo3();
+    PSVRuntimeInfo4 *pInfo4 = m_PSV.GetPSVRuntimeInfo4();
+
     if (pInfo)
       hlsl::SetShaderProps(pInfo, m_Module);
     if (pInfo1)
@@ -806,6 +808,8 @@ class DxilPSVWriter : public DxilPartWriter {
       hlsl::SetShaderProps(pInfo2, m_Module);
     if (pInfo3)
       pInfo3->EntryFunctionName = EntryFunctionName;
+    if (pInfo4)
+      hlsl::SetShaderProps(pInfo4, m_Module);
 
     // Set resource binding information
     UINT uResIndex = 0;

lib/DxilContainer/DxilPipelineStateValidation.cpp

@@ -33,6 +33,8 @@ uint32_t hlsl::GetPSVVersion(uint32_t ValMajor, uint32_t ValMinor) {
     PSVVersion = 1;
   else if (DXIL::CompareVersions(ValMajor, ValMinor, 1, 8) < 0)
     PSVVersion = 2;
+  else if (DXIL::CompareVersions(ValMajor, ValMinor, 1, 10) < 0)
+    PSVVersion = 3;
   return PSVVersion;
 }
 
@@ -305,6 +307,20 @@ void hlsl::SetShaderProps(PSVRuntimeInfo2 *pInfo2, const DxilModule &DM) {
   }
 }
 
+void hlsl::SetShaderProps(PSVRuntimeInfo4 *pInfo4, const DxilModule &DM) {
+  assert(pInfo4);
+  const ShaderModel *SM = DM.GetShaderModel();
+  switch (SM->GetKind()) {
+  case ShaderModel::Kind::Compute:
+  case ShaderModel::Kind::Mesh:
+  case ShaderModel::Kind::Amplification:
+    pInfo4->NumBytesGroupSharedMemory = DM.GetTGSMSizeInBytes();
+    break;
+  default:
+    break;
+  }
+}
+
 void PSVResourceBindInfo0::Print(raw_ostream &OS) const {
   OS << "PSVResourceBindInfo:\n";
   OS << "  Space: " << Space << "\n";
@@ -584,8 +600,9 @@ void PSVDependencyTable::Print(raw_ostream &OS, const char *InputSetName,
 
 void hlsl::PrintPSVRuntimeInfo(llvm::raw_ostream &OS, PSVRuntimeInfo0 *pInfo0,
                                PSVRuntimeInfo1 *pInfo1, PSVRuntimeInfo2 *pInfo2,
-                               PSVRuntimeInfo3 *pInfo3, uint8_t ShaderKind,
-                               const char *EntryName, const char *Comment) {
+                               PSVRuntimeInfo3 *pInfo3, PSVRuntimeInfo4 *pInfo4,
+                               uint8_t ShaderKind, const char *EntryName,
+                               const char *Comment) {
   if (pInfo1 && pInfo1->ShaderStage != ShaderKind)
     ShaderKind = pInfo1->ShaderStage;
   OS << Comment << "PSVRuntimeInfo:\n";
@@ -808,13 +825,23 @@ void hlsl::PrintPSVRuntimeInfo(llvm::raw_ostream &OS, PSVRuntimeInfo0 *pInfo0,
       OS << Comment << " NumThreads=(" << pInfo2->NumThreadsX << ","
          << pInfo2->NumThreadsY << "," << pInfo2->NumThreadsZ << ")\n";
     }
+    if (pInfo4) {
+      OS << Comment
+         << " NumBytesGroupSharedMemory: " << pInfo4->NumBytesGroupSharedMemory
+         << "\n";
+    }
     break;
   case PSVShaderKind::Amplification:
     OS << Comment << " Amplification Shader\n";
     if (pInfo2) {
       OS << Comment << " NumThreads=(" << pInfo2->NumThreadsX << ","
          << pInfo2->NumThreadsY << "," << pInfo2->NumThreadsZ << ")\n";
     }
+    if (pInfo4) {
+      OS << Comment
+         << " NumBytesGroupSharedMemory: " << pInfo4->NumBytesGroupSharedMemory
+         << "\n";
+    }
     break;
   case PSVShaderKind::Mesh:
     OS << Comment << " Mesh Shader\n";
@@ -841,6 +868,11 @@ void hlsl::PrintPSVRuntimeInfo(llvm::raw_ostream &OS, PSVRuntimeInfo0 *pInfo0,
       OS << Comment << " NumThreads=(" << pInfo2->NumThreadsX << ","
          << pInfo2->NumThreadsY << "," << pInfo2->NumThreadsZ << ")\n";
     }
+    if (pInfo4) {
+      OS << Comment
+         << " NumBytesGroupSharedMemory: " << pInfo4->NumBytesGroupSharedMemory
+         << "\n";
+    }
     break;
   case PSVShaderKind::Library:
   case PSVShaderKind::Invalid:
@@ -887,9 +919,10 @@ void DxilPipelineStateValidation::PrintPSVRuntimeInfo(
   PSVRuntimeInfo1 *pInfo1 = m_pPSVRuntimeInfo1;
   PSVRuntimeInfo2 *pInfo2 = m_pPSVRuntimeInfo2;
   PSVRuntimeInfo3 *pInfo3 = m_pPSVRuntimeInfo3;
+  PSVRuntimeInfo4 *pInfo4 = m_pPSVRuntimeInfo4;
 
   hlsl::PrintPSVRuntimeInfo(
-      OS, pInfo0, pInfo1, pInfo2, pInfo3, ShaderKind,
+      OS, pInfo0, pInfo1, pInfo2, pInfo3, pInfo4, ShaderKind,
       m_pPSVRuntimeInfo3 ? m_StringTable.Get(pInfo3->EntryFunctionName) : "",
       Comment);
 }

lib/DxilValidation/DxilContainerValidation.cpp

@@ -185,7 +185,8 @@ class PSVContentVerifier {
                            unsigned PSVVersion);
   void VerifyViewIDDependence(PSVRuntimeInfo1 *PSV1, unsigned PSVVersion);
   void VerifyEntryProperties(const ShaderModel *SM, PSVRuntimeInfo0 *PSV0,
-                             PSVRuntimeInfo1 *PSV1, PSVRuntimeInfo2 *PSV2);
+                             PSVRuntimeInfo1 *PSV1, PSVRuntimeInfo2 *PSV2,
+                             PSVRuntimeInfo3 *PSV3, PSVRuntimeInfo4 *PSV4);
   void EmitMismatchError(StringRef Name, StringRef PartContent,
                          StringRef ModuleContent) {
     ValCtx.EmitFormatError(ValidationRule::ContainerContentMatches,
@@ -409,16 +410,16 @@ void PSVContentVerifier::VerifyResources(unsigned PSVVersion) {
   VerifyResourceTable(DM.GetUAVs(), ResIndex, PSVVersion);
 }
 
-void PSVContentVerifier::VerifyEntryProperties(const ShaderModel *SM,
-                                               PSVRuntimeInfo0 *PSV0,
-                                               PSVRuntimeInfo1 *PSV1,
-                                               PSVRuntimeInfo2 *PSV2) {
-  PSVRuntimeInfo3 DMPSV;
-  memset(&DMPSV, 0, sizeof(PSVRuntimeInfo3));
+void PSVContentVerifier::VerifyEntryProperties(
+    const ShaderModel *SM, PSVRuntimeInfo0 *PSV0, PSVRuntimeInfo1 *PSV1,
+    PSVRuntimeInfo2 *PSV2, PSVRuntimeInfo3 *PSV3, PSVRuntimeInfo4 *PSV4) {
+  PSVRuntimeInfo4 DMPSV;
+  memset(&DMPSV, 0, sizeof(PSVRuntimeInfo4));
 
   hlsl::SetShaderProps((PSVRuntimeInfo0 *)&DMPSV, DM);
   hlsl::SetShaderProps((PSVRuntimeInfo1 *)&DMPSV, DM);
   hlsl::SetShaderProps((PSVRuntimeInfo2 *)&DMPSV, DM);
+  hlsl::SetShaderProps((PSVRuntimeInfo4 *)&DMPSV, DM);
   if (PSV1) {
     // Init things not set in InitPSVRuntimeInfo.
     DMPSV.ShaderStage = static_cast<uint8_t>(SM->GetKind());
@@ -444,10 +445,14 @@ void PSVContentVerifier::VerifyEntryProperties(const ShaderModel *SM,
   else
     Mismatched = memcmp(PSV0, &DMPSV, sizeof(PSVRuntimeInfo0)) != 0;
 
+  if (PSV4 &&
+      PSV4->NumBytesGroupSharedMemory != DMPSV.NumBytesGroupSharedMemory)
+    Mismatched = true;
+
   if (Mismatched) {
     std::string Str;
     raw_string_ostream OS(Str);
-    hlsl::PrintPSVRuntimeInfo(OS, &DMPSV, &DMPSV, &DMPSV, &DMPSV,
+    hlsl::PrintPSVRuntimeInfo(OS, &DMPSV, &DMPSV, &DMPSV, &DMPSV, &DMPSV,
                               static_cast<uint8_t>(SM->GetKind()),
                               DM.GetEntryFunctionName().c_str(), "");
     OS.flush();
@@ -476,9 +481,11 @@ void PSVContentVerifier::Verify(unsigned ValMajor, unsigned ValMinor,
   PSVRuntimeInfo0 *PSV0 = PSV.GetPSVRuntimeInfo0();
   PSVRuntimeInfo1 *PSV1 = PSV.GetPSVRuntimeInfo1();
   PSVRuntimeInfo2 *PSV2 = PSV.GetPSVRuntimeInfo2();
+  PSVRuntimeInfo3 *PSV3 = PSV.GetPSVRuntimeInfo3();
+  PSVRuntimeInfo4 *PSV4 = PSV.GetPSVRuntimeInfo4();
 
   const ShaderModel *SM = DM.GetShaderModel();
-  VerifyEntryProperties(SM, PSV0, PSV1, PSV2);
+  VerifyEntryProperties(SM, PSV0, PSV1, PSV2, PSV3, PSV4);
   if (PSVVersion > 0) {
     if (((PSV.GetSigInputElements() + PSV.GetSigOutputElements() +
           PSV.GetSigPatchConstOrPrimElements()) > 0) &&

lib/DxilValidation/DxilValidation.cpp

@@ -3935,6 +3935,18 @@ static void ValidateGlobalVariables(ValidationContext &ValCtx) {
     Rule = ValidationRule::SmMaxMSSMSize;
     MaxSize = DXIL::kMaxMSSMSize;
   }
+
+  // Check if the entry function has attribute to override TGSM size.
+  if (M.HasDxilEntryProps(M.GetEntryFunction())) {
+    DxilEntryProps &EntryProps = M.GetDxilEntryProps(M.GetEntryFunction());
+    if (EntryProps.props.IsCS()) {
+      unsigned SpecifiedTGSMSize = EntryProps.props.groupSharedLimitBytes;
+      if (SpecifiedTGSMSize > 0) {
+        MaxSize = SpecifiedTGSMSize;
+      }
+    }
+  }
+
   if (TGSMSize > MaxSize) {
     Module::global_iterator GI = M.GetModule()->global_end();
     GlobalVariable *GV = &*GI;