[llvm] [DXIL] Add support for root signature flag element in DXContainer (PR #123147)

[Justin Bogner via llvm-commits]

================
@@ -0,0 +1,229 @@
+//===- DXILRootSignature.cpp - DXIL 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 helper objects and APIs for working with DXIL
+///       Root Signatures.
+///
+//===----------------------------------------------------------------------===//
+#include "DXILRootSignature.h"
+#include "DirectX.h"
+#include "llvm/ADT/StringSwitch.h"
+#include "llvm/ADT/Twine.h"
+#include "llvm/Analysis/DXILMetadataAnalysis.h"
+#include "llvm/BinaryFormat/DXContainer.h"
+#include "llvm/IR/Constants.h"
+#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"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/raw_ostream.h"
+#include <cstdint>
+#include <optional>
+#include <utility>
+
+using namespace llvm;
+using namespace llvm::dxil;
+
+static bool reportError(LLVMContext *Ctx, Twine Message,
+                        DiagnosticSeverity Severity = DS_Error) {
+  Ctx->diagnose(DiagnosticInfoGeneric(Message, Severity));
+  return true;
+}
+
+static bool parseRootFlags(LLVMContext *Ctx, mcdxbc::RootSignatureDesc &RSD,
+                           MDNode *RootFlagNode) {
+
+  if (RootFlagNode->getNumOperands() != 2)
+    return reportError(Ctx, "Invalid format for RootFlag Element");
+
+  auto *Flag = mdconst::extract<ConstantInt>(RootFlagNode->getOperand(1));
+  RSD.Flags = Flag->getZExtValue();
+
+  return false;
+}
+
+static bool parseRootSignatureElement(LLVMContext *Ctx,
+                                      mcdxbc::RootSignatureDesc &RSD,
+                                      MDNode *Element) {
+  MDString *ElementText = cast<MDString>(Element->getOperand(0));
+  if (ElementText == nullptr)
+    return reportError(Ctx, "Invalid format for Root Element");
+
+  RootSignatureElementKind ElementKind =
+      StringSwitch<RootSignatureElementKind>(ElementText->getString())
+          .Case("RootFlags", RootSignatureElementKind::RootFlags)
+          .Default(RootSignatureElementKind::Error);
+
+  switch (ElementKind) {
+
+  case RootSignatureElementKind::RootFlags:
+    return parseRootFlags(Ctx, RSD, Element);
+  case RootSignatureElementKind::Error:
+    return reportError(Ctx, "Invalid Root Signature Element: " +
+                                ElementText->getString());
+  }
+
+  llvm_unreachable("Unhandled RootSignatureElementKind enum.");
+}
+
+static bool parse(LLVMContext *Ctx, mcdxbc::RootSignatureDesc &RSD,
+                  MDNode *Node) {
+  bool HasError = false;
+
+  // Loop through the Root Elements of the root signature.
+  for (const auto &Operand : Node->operands()) {
+    MDNode *Element = dyn_cast<MDNode>(Operand);
+    if (Element == nullptr)
+      return reportError(Ctx, "Missing Root Element Metadata Node.");
+
+    HasError = HasError || parseRootSignatureElement(Ctx, RSD, Element);
+  }
+
+  return HasError;
+}
+
+static bool validate(LLVMContext *Ctx, const mcdxbc::RootSignatureDesc &RSD) {
+  if (!dxbc::RootSignatureValidations::isValidRootFlag(RSD.Flags)) {
+    return reportError(Ctx, "Invalid Root Signature flag value");
+  }
+  return false;
+}
+
+static SmallDenseMap<const Function *, mcdxbc::RootSignatureDesc>
+analyzeModule(Module &M) {
+
+  /** Root Signature are specified as following in the metadata:
+
+    !dx.rootsignatures = !{!2} ; list of function/root signature pairs
+    !2 = !{ ptr @main, !3 } ; function, root signature
+    !3 = !{ !4, !5, !6, !7 } ; list of root signature elements
+
+    So for each MDNode inside dx.rootsignatures NamedMDNode
+    (the Root parameter of this function), the parsing process needs
+    to loop through each of its operands and process the function,
+    signature pair.
+ */
+
+  LLVMContext *Ctx = &M.getContext();
+
+  SmallDenseMap<const Function *, mcdxbc::RootSignatureDesc> RSDMap;
+
+  NamedMDNode *RootSignatureNode = M.getNamedMetadata("dx.rootsignatures");
+  if (RootSignatureNode == nullptr)
+    return RSDMap;
+
+  for (const auto &RSDefNode : RootSignatureNode->operands()) {
+    if (RSDefNode->getNumOperands() != 2) {
+      reportError(Ctx, "Invalid format for Root Signature Definition. Pairs "
+                       "of function, root signature expected.");
+      continue;
+    }
+
+    // Function was pruned during compilation.
+    const MDOperand &FunctionPointerMdNode = RSDefNode->getOperand(0);
+    if (FunctionPointerMdNode == nullptr) {
+      reportError(
+          Ctx, "Function associated with Root Signature definition is null.");
+      continue;
+    }
+
+    ValueAsMetadata *VAM =
+        llvm::dyn_cast<ValueAsMetadata>(FunctionPointerMdNode.get());
+    if (VAM == nullptr) {
+      reportError(Ctx, "First element of root signature is not a Value");
+      continue;
+    }
+
+    Function *F = dyn_cast<Function>(VAM->getValue());
+    if (F == nullptr) {
+      reportError(Ctx, "First element of root signature is not a Function");
+      continue;
+    }
+
+    MDNode *RootElementListNode =
+        dyn_cast<MDNode>(RSDefNode->getOperand(1).get());
+
+    if (RootElementListNode == nullptr) {
+      reportError(Ctx, "Missing Root Element List Metadata node.");
+    }
----------------
bogner wrote:

Should there be a `continue` here? It looks to me like we'll go ahead and use a nullptr later if this happens.

https://github.com/llvm/llvm-project/pull/123147