[DXIL] Add support for root signature flag element in DXContainer

llvm/lib/Target/DirectX/DXContainerGlobals.cpp

@@ -154,14 +154,16 @@ void DXContainerGlobals::addRootSignature(Module &M,
                                           SmallVector<GlobalValue *> &Globals) {
 
   auto &RSA = getAnalysis<RootSignatureAnalysisWrapper>();
+
   if (!RSA.getResult())
     return;
 
+  const ModuleRootSignature &MRS = RSA.getResult().value();
   SmallString<256> Data;
   raw_svector_ostream OS(Data);
 
   RootSignatureHeader RSH;
-  RSH.Flags = RSA.getResult()->Flags;
+  RSH.Flags = MRS.Flags;
 
   RSH.write(OS);
 

llvm/lib/Target/DirectX/DXILRootSignature.cpp

@@ -20,6 +20,7 @@
 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
@@ -56,16 +57,12 @@ static bool parseRootSignatureElement(LLVMContext *Ctx,
 
   RootSignatureElementKind ElementKind =
       StringSwitch<RootSignatureElementKind>(ElementText->getString())
-          .Case("RootFlags", RootSignatureElementKind::RootFlags)
-          .Default(RootSignatureElementKind::None);
+          .Case("RootFlags", RootSignatureElementKind::RootFlags);
 
   switch (ElementKind) {
 
   case RootSignatureElementKind::RootFlags:
     return parseRootFlags(Ctx, MRS, Element);
-  default:
-    return reportError(Ctx,
-                       "Invalid Root Element: " + ElementText->getString());
   }
 
   return true;
@@ -95,8 +92,14 @@ static bool parse(LLVMContext *Ctx, ModuleRootSignature *MRS, NamedMDNode *Root,
       continue;
     }
 
+    const MDOperand &FunctionPointerMdNode = Node->getOperand(0);
+    if (FunctionPointerMdNode == nullptr) {
+      // Function was pruned during compilation.
+      continue;
+    }
+
     ValueAsMetadata *VAM =
-        llvm::dyn_cast<ValueAsMetadata>(Node->getOperand(0).get());
+        llvm::dyn_cast<ValueAsMetadata>(FunctionPointerMdNode.get());
     if (VAM == nullptr) {
       HasError =
           reportError(Ctx, "First element of root signature is not a value");
@@ -140,24 +143,26 @@ static bool validate(LLVMContext *Ctx, ModuleRootSignature *MRS) {
   return false;
 }
 
-std::optional<ModuleRootSignature>
-ModuleRootSignature::analyzeModule(Module &M, ModuleMetadataInfo MMI) {
-  if (MMI.ShaderProfile == Triple::Library)
-    return std::nullopt;
+static const Function *getEntryFunction(Module &M, ModuleMetadataInfo MMI) {
 
   LLVMContext *Ctx = &M.getContext();
-
   if (MMI.EntryPropertyVec.size() != 1) {
     reportError(Ctx, "More than one entry function defined.");
-    return std::nullopt;
+    return nullptr;
   }
+  return MMI.EntryPropertyVec[0].Entry;
+}
+
+std::optional<ModuleRootSignature>
+ModuleRootSignature::analyzeModule(Module &M, const Function *F) {
+
+  LLVMContext *Ctx = &M.getContext();
 
   ModuleRootSignature MRS;
-  const Function *EntryFunction = MMI.EntryPropertyVec[0].Entry;
 
   NamedMDNode *RootSignatureNode = M.getNamedMetadata("dx.rootsignatures");
-  if (RootSignatureNode == nullptr ||
-      parse(Ctx, &MRS, RootSignatureNode, EntryFunction) || validate(Ctx, &MRS))
+  if (RootSignatureNode == nullptr || parse(Ctx, &MRS, RootSignatureNode, F) ||
+      validate(Ctx, &MRS))
     return std::nullopt;
 
   return MRS;
@@ -168,15 +173,18 @@ AnalysisKey RootSignatureAnalysis::Key;
 std::optional<ModuleRootSignature>
 RootSignatureAnalysis::run(Module &M, ModuleAnalysisManager &AM) {
   ModuleMetadataInfo MMI = AM.getResult<DXILMetadataAnalysis>(M);
-  return ModuleRootSignature::analyzeModule(M, MMI);
+  if (MMI.ShaderProfile == Triple::Library)
+    return std::nullopt;
+  return ModuleRootSignature::analyzeModule(M, getEntryFunction(M, MMI));
 }
 
 //===----------------------------------------------------------------------===//
 bool RootSignatureAnalysisWrapper::runOnModule(Module &M) {
-
   dxil::ModuleMetadataInfo &MMI =
       getAnalysis<DXILMetadataAnalysisWrapperPass>().getModuleMetadata();
-  MRS = ModuleRootSignature::analyzeModule(M, MMI);
+  if (MMI.ShaderProfile == Triple::Library)
+    return false;
+  MRS = ModuleRootSignature::analyzeModule(M, getEntryFunction(M, MMI));
   return false;
 }
 

llvm/lib/Target/DirectX/DXILRootSignature.h

@@ -22,12 +22,12 @@
 namespace llvm {
 namespace dxil {
 
-enum class RootSignatureElementKind { None = 0, RootFlags = 1 };
+enum class RootSignatureElementKind { RootFlags = 1 };
 
 struct ModuleRootSignature {
   uint32_t Flags = 0;
-  static std::optional<ModuleRootSignature>
-  analyzeModule(Module &M, ModuleMetadataInfo MMI);
+  static std::optional<ModuleRootSignature> analyzeModule(Module &M,
+                                                          const Function *F);
 };
 
 class RootSignatureAnalysis : public AnalysisInfoMixin<RootSignatureAnalysis> {

llvm/test/CodeGen/DirectX/ContainerData/RootSignature-MultipleEntryFunctions-Error.ll

@@ -0,0 +1,26 @@
+; RUN: not --crash llc %s --filetype=obj -o -
+; expected-error@-1: More than one entry function defined
+
+target triple = "dxil-unknown-shadermodel6.0-compute"
+
+
+define void @main() #0 {
+entry:
+  ret void
+}
+
+define void @anotherMain() #1 {
+entry:
+  ret void
+}
+
+attributes #0 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }
+attributes #1 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }
+
+!dx.rootsignatures = !{!2, !5} ; list of function/root signature pairs
+!2 = !{ ptr @main, !3 } ; function, root signature
+!3 = !{ !4 } ; list of root signature elements
+!4 = !{ !"RootFlags", i32 1 } ; 1 = allow_input_assembler_input_layout
+!5 = !{ ptr @anotherMain, !6 } ; function, root signature
+!6 = !{ !7 } ; list of root signature elements
+!7 = !{ !"RootFlags", i32 2 } ; 1 = allow_input_assembler_input_layout

llvm/test/CodeGen/DirectX/ContainerData/RootSignature-MultipleEntryFunctions.ll

@@ -0,0 +1,35 @@
+; RUN: llc %s --filetype=obj -o - | obj2yaml | FileCheck %s
+
+target triple = "dxil-unknown-shadermodel6.0-compute"
+
+
+define void @main() {
+entry:
+  ret void
+}
+
+define void @anotherMain() #0 {
+entry:
+  ret void
+}
+
+attributes #0 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }
+
+!dx.rootsignatures = !{!2, !5} ; list of function/root signature pairs
+!2 = !{ ptr @main, !3 } ; function, root signature
+!3 = !{ !4 } ; list of root signature elements
+!4 = !{ !"RootFlags", i32 1 } ; 1 = allow_input_assembler_input_layout
+!5 = !{ ptr @anotherMain, !6 } ; function, root signature
+!6 = !{ !7 } ; list of root signature elements
+!7 = !{ !"RootFlags", i32 2 } ; 1 = allow_input_assembler_input_layout
+
+
+; CHECK:  - Name:            RTS0
+; CHECK-NEXT:    Size:            24
+; CHECK-NEXT:    RootSignature:
+; CHECK-NEXT:      Version:         2
+; CHECK-NEXT:      NumParameters:   0
+; CHECK-NEXT:      RootParametersOffset: 0
+; CHECK-NEXT:      NumStaticSamplers: 0
+; CHECK-NEXT:      StaticSamplersOffset: 0
+; CHECK-NEXT:      DenyVertexShaderRootAccess: true