Add GroupSharedLimit as a shader attribute

Author: Jack Elliott

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 kDxilMaxGroupSharedMemTag = 24;
 
   // Node Input/Output State.
   static const unsigned kDxilNodeOutputIDTag = 0;

lib/DXIL/DxilMetadataHelper.cpp

@@ -1624,6 +1624,10 @@ MDTuple *DxilMDHelper::EmitDxilEntryProperties(uint64_t rawShaderFlag,
       }
       MDVals.emplace_back(MDNode::get(m_Ctx, WaveSizeVal));
     }
+
+    MDVals.emplace_back(Uint32ToConstMD(DxilMDHelper::kDxilMaxGroupSharedMemTag));
+    MDVals.emplace_back(
+        Uint32ToConstMD(props.groupSharedLimitBytes));
   } break;
     // Geometry shader.
   case DXIL::ShaderKind::Geometry: {
@@ -1773,6 +1777,11 @@ void DxilMDHelper::LoadDxilEntryProperties(const MDOperand &MDO,
       props.numThreads[2] = ConstMDToUint32(pNode->getOperand(2));
     } break;
 
+    case DxilMDHelper::kDxilMaxGroupSharedMemTag: {
+      DXASSERT(props.IsCS(), "else invalid shader kind");
+      props.groupSharedLimitBytes = ConstMDToUint32(MDO);
+    } break;
+
     case DxilMDHelper::kDxilGSStateTag: {
       DXASSERT(props.IsGS(), "else invalid shader kind");
       auto &GS = props.ShaderProps.GS;

lib/DxilValidation/DxilValidation.cpp

@@ -3921,6 +3921,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;

tools/clang/include/clang/Basic/Attr.td

@@ -671,6 +671,11 @@ def HLSLNumThreads: InheritableAttr {
   let Args = [IntArgument<"X">, IntArgument<"Y">, IntArgument<"Z">];
   let Documentation = [Undocumented];
 }
+def HLSLGroupSharedLimit: InheritableAttr {
+  let Spellings = [CXX11<"", "GroupSharedLimit", 2017>];
+  let Args = [IntArgument<"Limit">];
+  let Documentation = [Undocumented];
+}
 def HLSLRootSignature: InheritableAttr {
   let Spellings = [CXX11<"", "RootSignature", 2015>];
   let Args = [StringArgument<"SignatureName">];

tools/clang/lib/CodeGen/CGHLSLMS.cpp

@@ -1646,6 +1646,21 @@ void CGMSHLSLRuntime::AddHLSLFunctionInfo(Function *F, const FunctionDecl *FD) {
     }
   }
 
+  if (const HLSLGroupSharedLimitAttr *Attr = FD->getAttr<HLSLGroupSharedLimitAttr>()) {
+    funcProps->groupSharedLimitBytes = Attr->getLimit();
+
+    if (isEntry && !SM->IsCS() && !SM->IsMS() && !SM->IsAS()) {
+      unsigned DiagID = Diags.getCustomDiagID(
+          DiagnosticsEngine::Error,
+          "attribute GroupSharedLimit only valid for CS/MS/AS.");
+      Diags.Report(Attr->getLocation(), DiagID);
+      return;
+    }
+
+  } else {
+    funcProps->groupSharedLimitBytes = DXIL::kMaxTGSMSize;
+  }
+
   // Hull shader.
   if (const HLSLPatchConstantFuncAttr *Attr =
           FD->getAttr<HLSLPatchConstantFuncAttr>()) {

tools/clang/lib/Parse/ParseDecl.cpp

@@ -833,6 +833,7 @@ void Parser::ParseGNUAttributeArgs(IdentifierInfo *AttrName,
     case AttributeList::AT_HLSLMaxVertexCount:
     case AttributeList::AT_HLSLUnroll:
     case AttributeList::AT_HLSLWaveSize:
+    case AttributeList::AT_HLSLGroupSharedLimit:
     case AttributeList::AT_NoInline:
       // The following are not accepted in [attribute(param)] syntax:
       // case AttributeList::AT_HLSLCentroid:

tools/clang/lib/Sema/SemaHLSL.cpp

@@ -14656,6 +14656,11 @@ void hlsl::HandleDeclAttributeForHLSL(Sema &S, Decl *D, const AttributeList &A,
           S.Context.getAddrSpaceQualType(VD->getType(), DXIL::kTGSMAddrSpace));
     }
     break;
+  case AttributeList::AT_HLSLGroupSharedLimit:
+    declAttr = ::new (S.Context) HLSLGroupSharedLimitAttr(
+        A.getRange(), S.Context, ValidateAttributeIntArg(S, A),
+        A.getAttributeSpellingListIndex());
+    break;
   case AttributeList::AT_HLSLUniform:
     declAttr = ::new (S.Context) HLSLUniformAttr(
         A.getRange(), S.Context, A.getAttributeSpellingListIndex());