[HLSL][RootSignature] Metadata generation of RootFlags, RootConstants, RootDescriptors

clang/test/CodeGenHLSL/RootSignature.hlsl

@@ -1,16 +1,17 @@
 // RUN: %clang_cc1 -triple dxil-pc-shadermodel6.3-library -emit-llvm -o - %s | FileCheck %s
 
-// CHECK: !dx.rootsignatures = !{![[#FIRST_ENTRY:]], ![[#SECOND_ENTRY:]]}
+// CHECK: !dx.rootsignatures = !{![[#EMPTY_ENTRY:]], ![[#DT_ENTRY:]],
+// CHECK-SAME: ![[#RF_ENTRY:]], ![[#RC_ENTRY:]], ![[#RD_ENTRY:]]}
 
-// CHECK: ![[#FIRST_ENTRY]] = !{ptr @FirstEntry, ![[#EMPTY:]]}
+// CHECK: ![[#EMPTY_ENTRY]] = !{ptr @EmptyEntry, ![[#EMPTY:]]}
 // CHECK: ![[#EMPTY]] = !{}
 
 [shader("compute"), RootSignature("")]
 [numthreads(1,1,1)]
-void FirstEntry() {}
+void EmptyEntry() {}
 
-// CHECK: ![[#SECOND_ENTRY]] = !{ptr @SecondEntry, ![[#SECOND_RS:]]}
-// CHECK: ![[#SECOND_RS]] = !{![[#TABLE:]]}
+// CHECK: ![[#DT_ENTRY]] = !{ptr @DescriptorTableEntry, ![[#DT_RS:]]}
+// CHECK: ![[#DT_RS]] = !{![[#TABLE:]]}
 // CHECK: ![[#TABLE]] = !{!"DescriptorTable", i32 0, ![[#CBV:]], ![[#SRV:]]}
 // CHECK: ![[#CBV]] = !{!"CBV", i32 1, i32 0, i32 0, i32 -1, i32 4}
 // CHECK: ![[#SRV]] = !{!"SRV", i32 4, i32 42, i32 3, i32 32, i32 0}
@@ -22,10 +23,51 @@ void FirstEntry() {}
   ")"
 [shader("compute"), RootSignature(SampleDescriptorTable)]
 [numthreads(1,1,1)]
-void SecondEntry() {}
+void DescriptorTableEntry() {}
+
+// CHECK: ![[#RF_ENTRY]] = !{ptr @RootFlagsEntry, ![[#RF_RS:]]}
+// CHECK: ![[#RF_RS]] = !{![[#ROOT_FLAGS:]]}
+// CHECK: ![[#ROOT_FLAGS]] = !{!"RootFlags", i32 2114}
+
+#define SampleRootFlags \
+  "RootFlags( " \
+  " Deny_Vertex_Shader_Root_Access | Allow_Stream_Output | " \
+  " sampler_heap_directly_indexed " \
+  ")"
+[shader("compute"), RootSignature(SampleRootFlags)]
+[numthreads(1,1,1)]
+void RootFlagsEntry() {}
+
+// CHECK: ![[#RC_ENTRY]] = !{ptr @RootConstantsEntry, ![[#RC_RS:]]}
+// CHECK: ![[#RC_RS]] = !{![[#ROOT_CONSTANTS:]]}
+// CHECK: ![[#ROOT_CONSTANTS]] = !{!"RootConstants", i32 5, i32 1, i32 2, i32 1}
+
+#define SampleRootConstants \
+  "RootConstants(" \
+  " space = 2, " \
+  " visibility = Shader_Visibility_Pixel, " \
+  " b1, num32BitConstants = 1 " \
+  ")"
+[shader("compute"), RootSignature(SampleRootConstants)]
+[numthreads(1,1,1)]
+void RootConstantsEntry() {}
+
+// CHECK: ![[#RD_ENTRY]] = !{ptr @RootDescriptorsEntry, ![[#RD_RS:]]}
+// CHECK: ![[#RD_RS]] = !{![[#ROOT_CBV:]], ![[#ROOT_UAV:]], ![[#ROOT_SRV:]]}
+// CHECK: ![[#ROOT_CBV]] = !{!"RootCBV", i32 0, i32 0, i32 0, i32 4}
+// CHECK: ![[#ROOT_UAV]] = !{!"RootUAV", i32 0, i32 42, i32 3, i32 2}
+// CHECK: ![[#ROOT_SRV]] = !{!"RootSRV", i32 4, i32 0, i32 0, i32 2}
+
+#define SampleRootDescriptors \
+  "CBV(b0), " \
+  "UAV(space = 3, u42), " \
+  "SRV(t0, visibility = Shader_Visibility_Geometry, flags = Data_Volatile)"
+[shader("compute"), RootSignature(SampleRootDescriptors)]
+[numthreads(1,1,1)]
+void RootDescriptorsEntry() {}
 
 // Sanity test to ensure no root is added for this function as there is only
 // two entries in !dx.roosignatures
 [shader("compute")]
 [numthreads(1,1,1)]
-void ThirdEntry() {}
+void NoRSEntry() {}

llvm/include/llvm/Frontend/HLSL/HLSLRootSignature.h

@@ -188,6 +188,9 @@ class MetadataBuilder {
 
 private:
   /// Define the various builders for the different metadata types
+  MDNode *BuildRootFlags(const RootFlags &Flags);
+  MDNode *BuildRootConstants(const RootConstants &Constants);
+  MDNode *BuildRootDescriptor(const RootDescriptor &Descriptor);
   MDNode *BuildDescriptorTable(const DescriptorTable &Table);
   MDNode *BuildDescriptorTableClause(const DescriptorTableClause &Clause);
 

llvm/lib/Frontend/HLSL/HLSLRootSignature.cpp

@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Frontend/HLSL/HLSLRootSignature.h"
+#include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/bit.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Metadata.h"
@@ -171,9 +172,15 @@ void dumpRootElements(raw_ostream &OS, ArrayRef<RootElement> Elements) {
 MDNode *MetadataBuilder::BuildRootSignature() {
   for (const RootElement &Element : Elements) {
     MDNode *ElementMD = nullptr;
-    if (const auto &Clause = std::get_if<DescriptorTableClause>(&Element))
+    if (const auto &Flags = std::get_if<RootFlags>(&Element))
+      ElementMD = BuildRootFlags(*Flags);
+    else if (const auto &Constants = std::get_if<RootConstants>(&Element))
+      ElementMD = BuildRootConstants(*Constants);
+    else if (const auto &Descriptor = std::get_if<RootDescriptor>(&Element))
+      ElementMD = BuildRootDescriptor(*Descriptor);
+    else if (const auto &Clause = std::get_if<DescriptorTableClause>(&Element))
       ElementMD = BuildDescriptorTableClause(*Clause);
-    if (const auto &Table = std::get_if<DescriptorTable>(&Element))
+    else if (const auto &Table = std::get_if<DescriptorTable>(&Element))
       ElementMD = BuildDescriptorTable(*Table);
 
     // FIXME(#126586): remove once all RootElemnt variants are handled in a
@@ -187,6 +194,46 @@ MDNode *MetadataBuilder::BuildRootSignature() {
   return MDNode::get(Ctx, GeneratedMetadata);
 }
 
+MDNode *MetadataBuilder::BuildRootFlags(const RootFlags &Flags) {
+  IRBuilder<> Builder(Ctx);
+  Metadata *Operands[] = {
+      MDString::get(Ctx, "RootFlags"),
+      ConstantAsMetadata::get(Builder.getInt32(llvm::to_underlying(Flags))),
+  };
+  return MDNode::get(Ctx, Operands);
+}
+
+MDNode *MetadataBuilder::BuildRootConstants(const RootConstants &Constants) {
+  IRBuilder<> Builder(Ctx);
+  Metadata *Operands[] = {
+      MDString::get(Ctx, "RootConstants"),
+      ConstantAsMetadata::get(
+          Builder.getInt32(llvm::to_underlying(Constants.Visibility))),
+      ConstantAsMetadata::get(Builder.getInt32(Constants.Reg.Number)),
+      ConstantAsMetadata::get(Builder.getInt32(Constants.Space)),
+      ConstantAsMetadata::get(Builder.getInt32(Constants.Num32BitConstants)),
+  };
+  return MDNode::get(Ctx, Operands);
+}
+
+MDNode *MetadataBuilder::BuildRootDescriptor(const RootDescriptor &Descriptor) {
+  IRBuilder<> Builder(Ctx);
+  llvm::SmallString<7> Name;
+  llvm::raw_svector_ostream OS(Name);
+  OS << "Root" << ClauseType(llvm::to_underlying(Descriptor.Type));
+
+  Metadata *Operands[] = {
+      MDString::get(Ctx, OS.str()),
+      ConstantAsMetadata::get(
+          Builder.getInt32(llvm::to_underlying(Descriptor.Visibility))),
+      ConstantAsMetadata::get(Builder.getInt32(Descriptor.Reg.Number)),
+      ConstantAsMetadata::get(Builder.getInt32(Descriptor.Space)),
+      ConstantAsMetadata::get(
+          Builder.getInt32(llvm::to_underlying(Descriptor.Flags))),
+  };
+  return MDNode::get(Ctx, Operands);
+}
+
 MDNode *MetadataBuilder::BuildDescriptorTable(const DescriptorTable &Table) {
   IRBuilder<> Builder(Ctx);
   SmallVector<Metadata *> TableOperands;