adding metadata extraction

Author: joaosaffran

llvm/include/llvm/Analysis/DXILMetadataAnalysis.h

@@ -11,9 +11,11 @@
 
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/IR/PassManager.h"
+#include "llvm/MC/DXContainerRootSignature.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/VersionTuple.h"
 #include "llvm/TargetParser/Triple.h"
+#include <optional>
 
 namespace llvm {
 
@@ -37,6 +39,7 @@ struct ModuleMetadataInfo {
   Triple::EnvironmentType ShaderProfile{Triple::UnknownEnvironment};
   VersionTuple ValidatorVersion{};
   SmallVector<EntryProperties> EntryPropertyVec{};
+  std::optional<mcdxbc::RootSignatureDesc> RootSignatureDesc;
   void print(raw_ostream &OS) const;
 };
 

llvm/include/llvm/MC/DXContainerRootSignature.h

@@ -0,0 +1,74 @@
+//===- DXILRootSignature.h - 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 for working with DXIL Root
+/// Signatures.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_DIRECTX_HLSLROOTSIGNATURE_H
+#define LLVM_DIRECTX_HLSLROOTSIGNATURE_H
+
+#include "llvm/IR/Metadata.h"
+#include <memory>
+
+namespace llvm {
+namespace mcdxbc {
+
+enum class RootSignatureElementKind {
+  None = 0,
+  RootFlags = 1,
+  RootConstants = 2,
+  RootDescriptor = 3,
+  DescriptorTable = 4,
+  StaticSampler = 5
+};
+
+enum RootSignatureFlags : uint32_t {
+  None = 0,
+  AllowInputAssemblerInputLayout = 0x1,
+  DenyVertexShaderRootAccess = 0x2,
+  DenyHullShaderRootAccess = 0x4,
+  DenyDomainShaderRootAccess = 0x8,
+  DenyGeometryShaderRootAccess = 0x10,
+  DenyPixelShaderRootAccess = 0x20,
+  AllowStreamOutput = 0x40,
+  LocalRootSignature = 0x80,
+  DenyAmplificationShaderRootAccess = 0x100,
+  DenyMeshShaderRootAccess = 0x200,
+  CBVSRVUAVHeapDirectlyIndexed = 0x400,
+  SamplerHeapDirectlyIndexed = 0x800,
+  AllowLowTierReservedHwCbLimit = 0x80000000,
+  ValidFlags = 0x80000fff
+};
+
+struct RootSignatureDesc {
+  uint32_t Version;
+  RootSignatureFlags Flags;
+
+  void swapBytes() {
+    sys::swapByteOrder(Version);
+    sys::swapByteOrder(Flags);
+  }
+};
+
+class RootSignatureDescWriter {
+private:
+  RootSignatureDesc *Desc;
+
+public:
+  RootSignatureDescWriter(RootSignatureDesc *Desc) : Desc(Desc) {}
+
+  void write(raw_ostream &OS,
+             uint32_t Version = std::numeric_limits<uint32_t>::max());
+};
+
+} // namespace mcdxbc
+} // namespace llvm
+
+#endif // LLVM_DIRECTX_HLSLROOTSIGNATURE_H

llvm/lib/Analysis/DXILMetadataAnalysis.cpp

@@ -15,19 +15,99 @@
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/InitializePasses.h"
+#include "llvm/MC/DXContainerRootSignature.h"
+#include "llvm/Support/Casting.h"
 #include "llvm/Support/ErrorHandling.h"
+#include <memory>
 
 #define DEBUG_TYPE "dxil-metadata-analysis"
 
 using namespace llvm;
 using namespace dxil;
+using namespace llvm::mcdxbc;
+
+static bool parseRootFlags(MDNode *RootFlagNode, RootSignatureDesc *Desc) {
+
+  assert(RootFlagNode->getNumOperands() == 2 &&
+         "Invalid format for RootFlag Element");
+  auto *Flag = mdconst::extract<ConstantInt>(RootFlagNode->getOperand(1));
+  auto Value = (RootSignatureFlags)Flag->getZExtValue();
+
+  if ((Value & ~RootSignatureFlags::ValidFlags) != RootSignatureFlags::None)
+    return true;
+
+  Desc->Flags = Value;
+  return false;
+}
+
+static bool parseRootSignatureElement(MDNode *Element,
+                                      RootSignatureDesc *Desc) {
+  MDString *ElementText = cast<MDString>(Element->getOperand(0));
+
+  assert(ElementText != nullptr && "First preoperty of element is not ");
+
+  RootSignatureElementKind ElementKind =
+      StringSwitch<RootSignatureElementKind>(ElementText->getString())
+          .Case("RootFlags", RootSignatureElementKind::RootFlags)
+          .Case("RootConstants", RootSignatureElementKind::RootConstants)
+          .Case("RootCBV", RootSignatureElementKind::RootDescriptor)
+          .Case("RootSRV", RootSignatureElementKind::RootDescriptor)
+          .Case("RootUAV", RootSignatureElementKind::RootDescriptor)
+          .Case("Sampler", RootSignatureElementKind::RootDescriptor)
+          .Case("DescriptorTable", RootSignatureElementKind::DescriptorTable)
+          .Case("StaticSampler", RootSignatureElementKind::StaticSampler)
+          .Default(RootSignatureElementKind::None);
+
+  switch (ElementKind) {
+
+  case RootSignatureElementKind::RootFlags: {
+    return parseRootFlags(Element, Desc);
+    break;
+  }
+
+  case RootSignatureElementKind::RootConstants:
+  case RootSignatureElementKind::RootDescriptor:
+  case RootSignatureElementKind::DescriptorTable:
+  case RootSignatureElementKind::StaticSampler:
+  case RootSignatureElementKind::None:
+    llvm_unreachable("Not Implemented yet");
+    break;
+  }
+
+  return true;
+}
+
+bool parseRootSignature(RootSignatureDesc *Desc, int32_t Version,
+                        NamedMDNode *Root) {
+  Desc->Version = Version;
+  bool HasError = false;
+
+  for (unsigned int Sid = 0; Sid < Root->getNumOperands(); Sid++) {
+    // This should be an if, for error handling
+    MDNode *Node = cast<MDNode>(Root->getOperand(Sid));
+
+    // Not sure what use this for...
+    Metadata *Func = Node->getOperand(0).get();
+
+    // This should be an if, for error handling
+    MDNode *Elements = cast<MDNode>(Node->getOperand(1).get());
+
+    for (unsigned int Eid = 0; Eid < Elements->getNumOperands(); Eid++) {
+      MDNode *Element = cast<MDNode>(Elements->getOperand(Eid));
+
+      HasError = HasError || parseRootSignatureElement(Element, Desc);
+    }
+  }
+  return HasError;
+}
 
 static ModuleMetadataInfo collectMetadataInfo(Module &M) {
   ModuleMetadataInfo MMDAI;
   Triple TT(Triple(M.getTargetTriple()));
   MMDAI.DXILVersion = TT.getDXILVersion();
   MMDAI.ShaderModelVersion = TT.getOSVersion();
   MMDAI.ShaderProfile = TT.getEnvironment();
+
   NamedMDNode *ValidatorVerNode = M.getNamedMetadata("dx.valver");
   if (ValidatorVerNode) {
     auto *ValVerMD = cast<MDNode>(ValidatorVerNode->getOperand(0));
@@ -37,6 +117,15 @@ static ModuleMetadataInfo collectMetadataInfo(Module &M) {
         VersionTuple(MajorMD->getZExtValue(), MinorMD->getZExtValue());
   }
 
+  NamedMDNode *RootSignatureNode = M.getNamedMetadata("dx.rootsignatures");
+  if (RootSignatureNode) {
+    mcdxbc::RootSignatureDesc Desc;
+
+    parseRootSignature(&Desc, 1, RootSignatureNode);
+
+    MMDAI.RootSignatureDesc = Desc;
+  }
+
   // For all HLSL Shader functions
   for (auto &F : M.functions()) {
     if (!F.hasFnAttribute("hlsl.shader"))

llvm/lib/MC/CMakeLists.txt

@@ -1,6 +1,7 @@
 add_llvm_component_library(LLVMMC
   ConstantPools.cpp
   DXContainerPSVInfo.cpp
+  DXContainerRootSignature.cpp
   ELFObjectWriter.cpp
   GOFFObjectWriter.cpp
   MCAsmBackend.cpp

llvm/lib/MC/DXContainerRootSignature.cpp

@@ -0,0 +1,36 @@
+//===- DXContainerRootSignature.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 the parsing logic to extract root signature data
+///       from LLVM IR metadata.
+///
+//===----------------------------------------------------------------------===//
+
+#include "llvm/MC/DXContainerRootSignature.h"
+#include "llvm/ADT/StringSwitch.h"
+#include "llvm/BinaryFormat/DXContainer.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/Metadata.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/SwapByteOrder.h"
+#include <cassert>
+#include <cstdint>
+
+using namespace llvm;
+using namespace llvm::mcdxbc;
+
+void RootSignatureDescWriter::write(raw_ostream &OS, uint32_t Version) {
+  dxbc::RootSignatureDesc Out{Desc->Version, Desc->Flags};
+
+  if (sys::IsBigEndianHost) {
+    Out.swapBytes();
+  }
+
+  OS.write(reinterpret_cast<const char *>(&Out), sizeof(RootSignatureDesc));
+}

llvm/lib/Target/DirectX/DXContainerGlobals.cpp

@@ -23,9 +23,11 @@
 #include "llvm/IR/Module.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/MC/DXContainerPSVInfo.h"
+#include "llvm/MC/DXContainerRootSignature.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/MD5.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
+#include <optional>
 
 using namespace llvm;
 using namespace llvm::dxil;
@@ -41,6 +43,7 @@ class DXContainerGlobals : public llvm::ModulePass {
   GlobalVariable *buildSignature(Module &M, Signature &Sig, StringRef Name,
                                  StringRef SectionName);
   void addSignature(Module &M, SmallVector<GlobalValue *> &Globals);
+  void addRootSignature(Module &M, SmallVector<GlobalValue *> &Globals);
   void addResourcesForPSV(Module &M, PSVRuntimeInfo &PSV);
   void addPipelineStateValidationInfo(Module &M,
                                       SmallVector<GlobalValue *> &Globals);
@@ -73,6 +76,7 @@ bool DXContainerGlobals::runOnModule(Module &M) {
   Globals.push_back(getFeatureFlags(M));
   Globals.push_back(computeShaderHash(M));
   addSignature(M, Globals);
+  addRootSignature(M, Globals);
   addPipelineStateValidationInfo(M, Globals);
   appendToCompilerUsed(M, Globals);
   return true;
@@ -144,6 +148,26 @@ void DXContainerGlobals::addSignature(Module &M,
   Globals.emplace_back(buildSignature(M, OutputSig, "dx.osg1", "OSG1"));
 }
 
+void DXContainerGlobals::addRootSignature(Module &M,
+                                          SmallVector<GlobalValue *> &Globals) {
+
+  std::optional<RootSignatureDesc> Desc =
+      getAnalysis<DXILMetadataAnalysisWrapperPass>()
+          .getModuleMetadata()
+          .RootSignatureDesc;
+  if (!Desc.has_value())
+    return;
+
+  SmallString<256> Data;
+  raw_svector_ostream OS(Data);
+  RootSignatureDescWriter writer(&Desc.value());
+  writer.write(OS);
+
+  Constant *Constant =
+      ConstantDataArray::getString(M.getContext(), Data, /*AddNull*/ false);
+  Globals.emplace_back(buildContainerGlobal(M, Constant, "dx.rts0", "RTS0"));
+}
+
 void DXContainerGlobals::addResourcesForPSV(Module &M, PSVRuntimeInfo &PSV) {
   const DXILBindingMap &DBM =
       getAnalysis<DXILResourceBindingWrapperPass>().getBindingMap();