introduce a MetadataBuilder to handle the construction of nodes

Author: inbelic

clang/lib/CodeGen/CGHLSLRuntime.cpp

@@ -124,20 +124,14 @@ void addRootSignature(
     llvm::Function *Fn, llvm::Module &M) {
   auto &Ctx = M.getContext();
 
-  SmallVector<Metadata *> GeneratedMetadata;
-  for (auto Element : Elements) {
-    MDNode *ExampleRootElement = MDNode::get(Ctx, {});
-    GeneratedMetadata.push_back(ExampleRootElement);
-  }
-
-  MDNode *ExampleRootSignature = MDNode::get(Ctx, GeneratedMetadata);
-
-  MDNode *ExamplePairing = MDNode::get(Ctx, {ValueAsMetadata::get(Fn),
-                                             ExampleRootSignature});
+  llvm::hlsl::root_signature::MetadataBuilder Builder(Ctx, Elements);
+  MDNode *RootSignature = Builder.BuildRootSignature();
+  MDNode *FnPairing = MDNode::get(Ctx, {ValueAsMetadata::get(Fn),
+                                        RootSignature});
 
   StringRef RootSignatureValKey = "dx.rootsignatures";
   auto *RootSignatureValMD = M.getOrInsertNamedMetadata(RootSignatureValKey);
-  RootSignatureValMD->addOperand(ExamplePairing);
+  RootSignatureValMD->addOperand(FnPairing);
 }
 
 } // namespace

clang/test/CodeGenHLSL/RootSignature.hlsl

@@ -1,20 +1,25 @@
 // RUN: %clang_cc1 -triple dxil-pc-shadermodel6.3-library -emit-llvm -o - %s | FileCheck %s
 
 // CHECK-DAG: ![[#EMPTY:]] = !{}
-// CHECK-DAG: ![[#SECOND_RS:]] = !{![[#EMPTY]]}
-// CHECK-DAG: ![[#SECOND_ENTRY:]] = !{ptr @SecondEntry, ![[#SECOND_RS]]}
-// CHECK-DAG: ![[#FIRST_ENTRY:]] = !{ptr @FirstEntry, ![[#EMPTY]]}
-// CHECK-DAG: !dx.rootsignatures = !{![[#FIRST_ENTRY]], ![[#SECOND_ENTRY]]}
-
 [shader("compute"), RootSignature("")]
 [numthreads(1,1,1)]
 void FirstEntry() {}
 
-[shader("compute"), RootSignature("DescriptorTable()")]
+// CHECK-DAG: ![[#TABLE:]] = !{!"DescriptorTable"}
+// CHECK-DAG: ![[#SECOND_RS:]] = !{![[#TABLE]]}
+
+#define SampleDescriptorTable \
+  "DescriptorTable( " \
+  ")"
+[shader("compute"), RootSignature(SampleDescriptorTable)]
 [numthreads(1,1,1)]
 void SecondEntry() {}
 
 // Sanity test to ensure to root is added for this function
 [shader("compute")]
 [numthreads(1,1,1)]
 void ThirdEntry() {}
+
+// CHECK-DAG: ![[#FIRST_ENTRY:]] = !{ptr @FirstEntry, ![[#EMPTY]]}
+// CHECK-DAG: ![[#SECOND_ENTRY:]] = !{ptr @SecondEntry, ![[#SECOND_RS]]}
+// CHECK-DAG: !dx.rootsignatures = !{![[#FIRST_ENTRY]], ![[#SECOND_ENTRY]]}

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

@@ -14,10 +14,16 @@
 #ifndef LLVM_FRONTEND_HLSL_HLSLROOTSIGNATURE_H
 #define LLVM_FRONTEND_HLSL_HLSLROOTSIGNATURE_H
 
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLForwardCompat.h"
 #include "llvm/Support/DXILABI.h"
 #include <variant>
 
 namespace llvm {
+class LLVMContext;
+class MDNode;
+class Metadata;
+
 namespace hlsl {
 namespace rootsig {
 
@@ -122,6 +128,24 @@ using RootElement = std::variant<DescriptorTable, DescriptorTableClause>;
 using ParamType = std::variant<uint32_t *, DescriptorRangeOffset *,
                                DescriptorRangeFlags *, ShaderVisibility *>;
 
+class MetadataBuilder {
+ public:
+  MetadataBuilder(llvm::LLVMContext &Ctx, ArrayRef<RootElement> Elements)
+   : Ctx(Ctx), Elements(Elements) {}
+
+  // Iterates through the elements and builds the respective nodes
+  MDNode *BuildRootSignature();
+
+ private:
+  // Define the various builders for the different metadata types
+  MDNode *BuildDescriptorTable(const DescriptorTable &Table);
+  MDNode *BuildDescriptorTableClause(const DescriptorTableClause &Clause);
+
+  llvm::LLVMContext &Ctx;
+  ArrayRef<RootElement> Elements;
+  SmallVector<Metadata *> GeneratedMetadata;
+};
+
 } // namespace rootsig
 } // namespace hlsl
 } // namespace llvm

llvm/lib/Frontend/HLSL/CMakeLists.txt

@@ -1,5 +1,6 @@
 add_llvm_component_library(LLVMFrontendHLSL
   HLSLResource.cpp
+  HLSLRootSignature.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${LLVM_MAIN_INCLUDE_DIR}/llvm/Frontend

llvm/lib/Frontend/HLSL/HLSLRootSignature.cpp

@@ -0,0 +1,77 @@
+//===- HLSLRootSignature.cpp - HLSL Root Signature helper objects ----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file This file contains helpers for working with HLSL Root Signatures.
+///
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Frontend/HLSL/HLSLRootSignature.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Metadata.h"
+#include "llvm/IR/Module.h"
+
+namespace llvm {
+namespace hlsl {
+namespace root_signature {
+
+static MDString *ClauseTypeToName(LLVMContext &Ctx, ClauseType Type) {
+  StringRef Name;
+  switch (Type) {
+  case ClauseType::CBuffer:
+    Name = "CBV";
+    break;
+  case ClauseType::SRV:
+    Name = "SRV";
+    break;
+  case ClauseType::UAV:
+    Name = "UAV";
+    break;
+  case ClauseType::Sampler:
+    Name = "Sampler";
+    break;
+  }
+  return MDString::get(Ctx, Name);
+}
+
+// Helper struct so that we can use the overloaded notation of std::visit
+template <class... Ts> struct OverloadBuilds : Ts... {
+  using Ts::operator()...;
+};
+template <class... Ts> OverloadBuilds(Ts...) -> OverloadBuilds<Ts...>;
+
+MDNode *MetadataBuilder::BuildRootSignature() {
+  for (const RootElement &Element : Elements) {
+    MDNode *ElementMD =
+      std::visit(
+        OverloadBuilds{
+            [&](DescriptorTable Table) -> MDNode * {
+              return BuildDescriptorTable(Table);
+            },
+            [&](DescriptorTableClause Clause) -> MDNode * {
+              return BuildDescriptorTableClause(Clause);
+            },
+        },
+        Element);
+    GeneratedMetadata.push_back(ElementMD);
+  }
+
+  return MDNode::get(Ctx, GeneratedMetadata);
+}
+
+MDNode *MetadataBuilder::BuildDescriptorTable(const DescriptorTable &Table) {
+  return MDNode::get(Ctx, {MDString::get(Ctx, "DescriptorTable")});
+}
+
+MDNode *MetadataBuilder::BuildDescriptorTableClause(const DescriptorTableClause &Clause) {
+  return MDNode::get(Ctx, {ClauseTypeToName(Ctx, Clause.Type)});
+}
+
+} // namespace root_signature
+} // namespace hlsl
+} // namespace llvm
+