[DirectX backend] add pass to collect Resource information from dxil metadata.

[python3kgae]

Summary

This is for container builder which need resource information after dxil metadata generation.

New analysis pass DXILMetadataResourceAnalysis is added to collect the information.
New constructor for dxil::ResourceBase/UAV/ConstantBuffer is added to build from dxil metadata.

[beanz]

I don't think this PR is correct. It seems like we should be able to ensure that the analysis result remains preserved after the metadata is generated allowing us to re-use the analysis result that was used to generate the metadata.

llvm/lib/Target/DirectX/DXILResource.cpp

45

If the tuple _must_ have the operand (as demonstrated by the asserts below), this can just be a dyn_cast.

dyn_cast_or_null has been replaced by dyn_cast_if_present, and should only be used when null is a valid input and output.

[python3kgae]

Add

 void getAnalysisUsage(AnalysisUsage &AU) const override {
  AU.addPreserved<DXILResourceWrapper>();
}

to DXILPrepareModule resolved the issue.

llvm/lib/Target/DirectX/CMakeLists.txt

add_llvm_target(DirectXCodeGen
   CBufferDataLayout.cpp
   DirectXAsmPrinter.cpp
   DirectXInstrInfo.cpp
   DirectXRegisterInfo.cpp
   DirectXSubtarget.cpp
   DirectXTargetMachine.cpp
   DXContainerGlobals.cpp
   DXILMetadata.cpp
+  DXILMetadataResourceAnalysis.cpp
   DXILOpBuilder.cpp
   DXILOpLowering.cpp
   DXILPrepare.cpp
   DXILPrettyPrinter.cpp
   DXILResource.cpp
   DXILResourceAnalysis.cpp
   DXILShaderFlags.cpp
   DXILTranslateMetadata.cpp

llvm/lib/Target/DirectX/DXILMetadataResourceAnalysis.h

+//===- DXILMetadataResourceAnalysis.h   - DXIL Metadata Resource analysis--===//
+//
+// 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 Analysis which collect DXIL resources from DXIL
+/// metadata.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIB_TARGET_DIRECTX_DXILMETADATARESOURCEANALYSIS_H
+#define LLVM_LIB_TARGET_DIRECTX_DXILMETADATARESOURCEANALYSIS_H
+
+#include "DXILResource.h"
+#include "llvm/IR/PassManager.h"
+#include "llvm/Pass.h"
+#include <memory>
+
+namespace llvm {
+/// Analysis pass that exposes the \c DXILResource for a module from dxil
+/// metadata.
+class DXILMetadataResourceAnalysis
+    : public AnalysisInfoMixin<DXILMetadataResourceAnalysis> {
+  friend AnalysisInfoMixin<DXILMetadataResourceAnalysis>;
+  static AnalysisKey Key;
+
+public:
+  typedef dxil::Resources Result;
+  dxil::Resources run(Module &M, ModuleAnalysisManager &AM);
+};
+
+/// Printer pass for the \c DXILMetadataResourceAnalysis results.
+class DXILMetadataResourcePrinterPass
+    : public PassInfoMixin<DXILMetadataResourcePrinterPass> {
+  raw_ostream &OS;
+
+public:
+  explicit DXILMetadataResourcePrinterPass(raw_ostream &OS) : OS(OS) {}
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
+};
+
+/// The legacy pass manager's analysis pass to compute DXIL resource
+/// information from DXIL metadata.
+class DXILMetadataResourceWrapper : public ModulePass {
+  dxil::Resources Resources;
+
+public:
+  static char ID; // Pass identification, replacement for typeid
+
+  DXILMetadataResourceWrapper();
+
+  dxil::Resources &getDXILResource() { return Resources; }
+  const dxil::Resources &getDXILResource() const { return Resources; }
+
+  /// Calculate the DXILResource for the module.
+  bool runOnModule(Module &M) override;
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.setPreservesAll();
+  }
+  void print(raw_ostream &O, const Module *M = nullptr) const override;
+};
+} // namespace llvm
+
+#endif

llvm/lib/Target/DirectX/DXILMetadataResourceAnalysis.cpp

+//===- DXILMetadataResourceAnalysis.cpp - DXIL Metadata Resource analysis--===//
+//
+// 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 Analysis which collect DXIL resources from DXIL
+/// metadata.
+///
+//===----------------------------------------------------------------------===//
+
+#include "DXILMetadataResourceAnalysis.h"
+#include "DirectX.h"
+#include "llvm/IR/PassManager.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "dxil-metadata-resource-analysis"
+
+dxil::Resources DXILMetadataResourceAnalysis::run(Module &M,
+                                                  ModuleAnalysisManager &AM) {
+  dxil::Resources R;
+  R.collectFromDXILMetadata(M);
+  return R;
+}
+
+AnalysisKey DXILMetadataResourceAnalysis::Key;
+
+PreservedAnalyses
+DXILMetadataResourcePrinterPass::run(Module &M, ModuleAnalysisManager &AM) {
+  dxil::Resources Res = AM.getResult<DXILMetadataResourceAnalysis>(M);
+  Res.print(OS);
+  return PreservedAnalyses::all();
+}
+
+char DXILMetadataResourceWrapper::ID = 0;
+INITIALIZE_PASS_BEGIN(DXILMetadataResourceWrapper, DEBUG_TYPE,
+                      "DXIL resource Information from metadata", true, true)
+INITIALIZE_PASS_END(DXILMetadataResourceWrapper, DEBUG_TYPE,
+                    "DXIL resource Information from metadata", true, true)
+
+bool DXILMetadataResourceWrapper::runOnModule(Module &M) {
+  Resources.collectFromDXILMetadata(M);
+  return false;
+}
+
+DXILMetadataResourceWrapper::DXILMetadataResourceWrapper() : ModulePass(ID) {}
+
+void DXILMetadataResourceWrapper::print(raw_ostream &OS, const Module *) const {
+  Resources.print(OS);
+}

llvm/lib/Target/DirectX/DXILResource.h

 #include "llvm/Frontend/HLSL/HLSLResource.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/Support/Compiler.h"
 #include <cstdint>
 
 namespace llvm {
 class Module;
 class GlobalVariable;
+class MDTuple;
+class MDOperand;
 
 namespace dxil {
 class CBufferDataLayout;
 
 class ResourceBase {
 protected:
   uint32_t ID;
   GlobalVariable *GV;
   StringRef Name;
   uint32_t Space;
   uint32_t LowerBound;
   uint32_t RangeSize;
   ResourceBase(uint32_t I, hlsl::FrontendResource R);
+  ResourceBase(const MDTuple *ResMD);
 
   void write(LLVMContext &Ctx, MutableArrayRef<Metadata *> Entries) const;
 
   void print(raw_ostream &O, StringRef IDPrefix, StringRef BindingPrefix) const;
   using Kinds = hlsl::ResourceKind;
   static StringRef getKindName(Kinds Kind);
   static void printKind(Kinds Kind, unsigned Alignment, raw_ostream &OS,
                         bool SRV = false, bool HasCounter = false,
struct ExtendedProperties {
     enum Tags : uint32_t {
       TypedBufferElementType = 0,
       StructuredBufferElementStride,
       SamplerFeedbackKind,
       Atomic64Use
     };
 
     MDNode *write(LLVMContext &Ctx) const;
+    void loadExtProps(const MDTuple *ExtTuple);
   };
 };
 
 class UAVResource : public ResourceBase {
   ResourceBase::Kinds Shape;
   bool GloballyCoherent;
   bool HasCounter;
   bool IsROV;
   ResourceBase::ExtendedProperties ExtProps;
 
   void parseSourceType(StringRef S);
 
 public:
   UAVResource(uint32_t I, hlsl::FrontendResource R);
+  UAVResource(const MDTuple *ResMD);
 
   MDNode *write() const;
   void print(raw_ostream &O) const;
 };
 
 class ConstantBuffer : public ResourceBase {
   uint32_t CBufferSizeInBytes = 0; // Cbuffer used size in bytes.
 public:
   ConstantBuffer(uint32_t I, hlsl::FrontendResource R);
+  ConstantBuffer(const MDTuple *ResMD);
+
   void setSize(CBufferDataLayout &DL);
   MDNode *write() const;
   void print(raw_ostream &O) const;
 };
 
 template <typename T> class ResourceTable {
   StringRef MDName;
 
   llvm::SmallVector<T> Data;
 
 public:
   ResourceTable(StringRef Name) : MDName(Name) {}
   void collect(Module &M);
+  void collectFromDXILMetadata(const MDOperand &MDO);
   MDNode *write(Module &M) const;
   void print(raw_ostream &O) const;
 };
 
 // FIXME: Fully computing the resource structures requires analyzing the IR
 // because some flags are set based on what operations are performed on the
 // resource. This partial patch handles some of the leg work, but not all of it.
 // See issue https://github.com/llvm/llvm-project/issues/57936.
 class Resources {
   ResourceTable<UAVResource> UAVs = {"hlsl.uavs"};
   ResourceTable<ConstantBuffer> CBuffers = {"hlsl.cbufs"};
 
 public:
   void collect(Module &M);
+  void collectFromDXILMetadata(Module &M);
   void write(Module &M) const;
   void print(raw_ostream &O) const;
   LLVM_DUMP_METHOD void dump() const;
 };
 
 } // namespace dxil
 } // namespace llvm
 
 #endif // LLVM_TARGET_DIRECTX_DXILRESOURCE_H

llvm/lib/Target/DirectX/DXILResource.cpp

if (!Entry || Entry->getNumOperands() == 0)
     return;
 
   uint32_t Counter = 0;
   for (auto *Res : Entry->operands()) {
     Data.push_back(T(Counter++, FrontendResource(cast<MDNode>(Res))));
   }
 }
 
+template <typename T>
+void ResourceTable<T>::collectFromDXILMetadata(const MDOperand &MDO) {
+  const MDTuple *Tuple = dyn_cast_or_null<MDTuple>(MDO.get());
+  if (!Tuple)
+    return;
+
+  for (unsigned i = 0; i < Tuple->getNumOperands(); i++) {
+    const MDTuple *ResMD =
+        dyn_cast_or_null<MDTuple>(Tuple->getOperand(i).get());
+    assert(ResMD != nullptr && "");
+    assert(ResMD->getNumOperands() >= 6 && "");
+    Data.emplace_back(T(ResMD));
+  }
+}
+
 template <> void ResourceTable<ConstantBuffer>::collect(Module &M) {
   NamedMDNode *Entry = M.getNamedMetadata(MDName);
   if (!Entry || Entry->getNumOperands() == 0)
     return;
 
   uint32_t Counter = 0;
   for (auto *Res : Entry->operands()) {
     Data.push_back(
         ConstantBuffer(Counter++, FrontendResource(cast<MDNode>(Res))));
   }
   // FIXME: share CBufferDataLayout with CBuffer load lowering.
   //   See https://github.com/llvm/llvm-project/issues/58381
   CBufferDataLayout CBDL(M.getDataLayout(), /*IsLegacy*/ true);
   for (auto &CB : Data)
     CB.setSize(CBDL);
 }
 
 void Resources::collect(Module &M) {
   UAVs.collect(M);
   CBuffers.collect(M);
 }
 
+void Resources::collectFromDXILMetadata(Module &M) {
+  NamedMDNode *ResNamedMD = M.getNamedMetadata("dx.resources");
+  if (!ResNamedMD)
+    return;
+  MDNode *ResTab = ResNamedMD->getOperand(0);
+  if (ResTab->getNumOperands() != 4)
+    return;
+
+  UAVs.collectFromDXILMetadata(
+      ResTab->getOperand(static_cast<int>(hlsl::ResourceClass::UAV)));
+  CBuffers.collectFromDXILMetadata(
+      ResTab->getOperand(static_cast<int>(hlsl::ResourceClass::CBuffer)));
+}
+
 ResourceBase::ResourceBase(uint32_t I, FrontendResource R)
     : ID(I), GV(R.getGlobalVariable()), Name(""), Space(R.getSpace()),
       LowerBound(R.getResourceIndex()), RangeSize(1) {
   if (auto *ArrTy = dyn_cast<ArrayType>(GV->getValueType()))
     RangeSize = ArrTy->getNumElements();
 }
 
+ResourceBase::ResourceBase(const MDTuple *ResMD) {
+  ID =
+      mdconst::extract<ConstantInt>(ResMD->getOperand(0))->getZExtValue();
+  Constant *GlobalSymbol = mdconst::extract<Constant>(
+      ResMD->getOperand(1));
+  GV = cast<GlobalVariable>(GlobalSymbol);
+  Name = dyn_cast<MDString>(ResMD->getOperand(2).get())->getString();
+  Space =
+      mdconst::extract<ConstantInt>(ResMD->getOperand(3))->getZExtValue();
+  LowerBound =
+      mdconst::extract<ConstantInt>(ResMD->getOperand(4))->getZExtValue();
+  RangeSize =
+      mdconst::extract<ConstantInt>(ResMD->getOperand(5))->getZExtValue();
+}
+
 StringRef ResourceBase::getComponentTypeName(ComponentType CompType) {
   switch (CompType) {
   case ComponentType::LastEntry:
   case ComponentType::Invalid:
     return "invalid";
   case ComponentType::I1:
     return "i1";
   case ComponentType::I16:

 
 UAVResource::UAVResource(uint32_t I, FrontendResource R)
     : ResourceBase(I, R),
       Shape(static_cast<ResourceBase::Kinds>(R.getResourceKind())),
       GloballyCoherent(false), HasCounter(false), IsROV(false), ExtProps() {
   parseSourceType(R.getSourceType());
 }
 
+UAVResource::UAVResource(const MDTuple *ResMD) : ResourceBase(ResMD) {
+  // UAV-specific fields.
+  Shape = static_cast<ResourceBase::Kinds>(
+      mdconst::extract<ConstantInt>(ResMD->getOperand(6))->getZExtValue());
+  GloballyCoherent =
+      mdconst::extract<ConstantInt>(ResMD->getOperand(7))->getZExtValue();
+  HasCounter =
+      mdconst::extract<ConstantInt>(ResMD->getOperand(8))->getZExtValue();
+  IsROV = mdconst::extract<ConstantInt>(ResMD->getOperand(9))->getZExtValue();
+  if (const MDTuple *ExtTuple = dyn_cast_or_null<MDTuple>(ResMD->getOperand(10).get()))
+    ExtProps.loadExtProps(ExtTuple);
+}
+
 void UAVResource::print(raw_ostream &OS) const {
   OS << "; " << left_justify(Name, 31);
 
   OS << right_justify("UAV", 10);
 
   printComponentType(
       Shape, ExtProps.ElementType.value_or(ComponentType::Invalid), 8, OS);
 
ComponentType ElTy = StringSwitch<ResourceBase::ComponentType>(S)
                            .Default(ComponentType::Invalid);
   if (ElTy != ComponentType::Invalid)
     ExtProps.ElementType = ElTy;
 }
 
 ConstantBuffer::ConstantBuffer(uint32_t I, hlsl::FrontendResource R)
     : ResourceBase(I, R) {}
 
+ConstantBuffer::ConstantBuffer(const MDTuple *ResMD) : ResourceBase(ResMD) {
+  CBufferSizeInBytes =
+      mdconst::extract<ConstantInt>(ResMD->getOperand(6))->getZExtValue();
+}
+
 void ConstantBuffer::setSize(CBufferDataLayout &DL) {
   CBufferSizeInBytes = DL.getTypeAllocSizeInBytes(GV->getValueType());
 }
 
 void ConstantBuffer::print(raw_ostream &OS) const {
   OS << "; " << left_justify(Name, 31);
 
   OS << right_justify("cbuffer", 10);
 
   printComponentType(Kinds::CBuffer, ComponentType::Invalid, 8, OS);
 
   printKind(Kinds::CBuffer, 12, OS, /*SRV*/ false, /*HasCounter*/ false);
   // Print the binding part.
   ResourceBase::print(OS, "CB", "cb");
 }
 
 template <typename T> void ResourceTable<T>::print(raw_ostream &OS) const {
   for (auto &Res : Data)
     Res.print(OS);
 }
 
+void ResourceBase::ExtendedProperties::loadExtProps(const MDTuple *ExtTuple) {
+  for (unsigned i = 0; i < ExtTuple->getNumOperands(); i += 2) {
+    Tags ExtTag = static_cast<Tags>(mdconst::extract<ConstantInt>(ExtTuple->getOperand(0))->getZExtValue());
+    const MDOperand &MDO = ExtTuple->getOperand(i + 1);
+    switch (ExtTag) {
+      case Tags::TypedBufferElementType:
+        ElementType = static_cast<ComponentType>(mdconst::extract<ConstantInt>(MDO)->getZExtValue());
+        break;
+      case Tags::StructuredBufferElementStride:
+      case Tags::SamplerFeedbackKind:
+      case Tags::Atomic64Use:
+        // FIXME: support all ext props.
+        assert(0 && "not supported ext props");
+        break;
+    }
+  }
+}
 MDNode *ResourceBase::ExtendedProperties::write(LLVMContext &Ctx) const {
   IRBuilder<> B(Ctx);
   SmallVector<Metadata *> Entries;
   if (ElementType) {
     Entries.emplace_back(
         ConstantAsMetadata::get(B.getInt32(TypedBufferElementType)));
     Entries.emplace_back(ConstantAsMetadata::get(
         B.getInt32(static_cast<uint32_t>(*ElementType))));

llvm/lib/Target/DirectX/DirectX.h

 ModulePass *createDXILOpLoweringLegacyPass();
 
 /// Initializer for DXILTranslateMetadata.
 void initializeDXILTranslateMetadataPass(PassRegistry &);
 
 /// Pass to emit metadata for DXIL.
 ModulePass *createDXILTranslateMetadataPass();
 
-/// Initializer for DXILTranslateMetadata.
+/// Initializer for DXILResourceWrapper.
 void initializeDXILResourceWrapperPass(PassRegistry &);
 
+/// Initializer for DXILMetadataResourceWrapper.
+void initializeDXILMetadataResourceWrapperPass(PassRegistry &);
+
 /// Pass to pretty print DXIL metadata.
 ModulePass *createDXILPrettyPrinterPass(raw_ostream &OS);
 
 /// Initializer for DXILPrettyPrinter.
 void initializeDXILPrettyPrinterPass(PassRegistry &);
 
 /// Initializer for dxil::ShaderFlagsAnalysisWrapper pass.
 void initializeShaderFlagsAnalysisWrapperPass(PassRegistry &);

llvm/lib/Target/DirectX/DirectXTargetMachine.cpp

 //===- DirectXTargetMachine.cpp - DirectX Target Implementation -*- C++ -*-===//
 //
 // 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 DirectX target initializer.
 ///
 //===----------------------------------------------------------------------===//
 
 #include "DirectXTargetMachine.h"
+#include "DXILMetadataResourceAnalysis.h"
 #include "DXILResourceAnalysis.h"
 #include "DXILShaderFlags.h"
 #include "DXILWriter/DXILWriterPass.h"
 #include "DirectX.h"
 #include "DirectXSubtarget.h"
 #include "DirectXTargetTransformInfo.h"
 #include "TargetInfo/DirectXTargetInfo.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"

 void DirectXTargetMachine::registerPassBuilderCallbacks(PassBuilder &PB) {
   PB.registerPipelineParsingCallback(
       [](StringRef PassName, ModulePassManager &PM,
          ArrayRef<PassBuilder::PipelineElement>) {
         if (PassName == "print-dxil-resource") {
           PM.addPass(DXILResourcePrinterPass(dbgs()));
           return true;
         }
+        if (PassName == "print-dxil-resource-from-dxil-metadata") {
+          PM.addPass(DXILMetadataResourcePrinterPass(dbgs()));
+          return true;
+        }
         if (PassName == "print-dx-shader-flags") {
           PM.addPass(dxil::ShaderFlagsAnalysisPrinter(dbgs()));
           return true;
         }
         return false;
       });
 
   PB.registerAnalysisRegistrationCallback([](ModuleAnalysisManager &MAM) {
     MAM.registerPass([&] { return DXILResourceAnalysis(); });
+    MAM.registerPass([&] { return DXILMetadataResourceAnalysis(); });
     MAM.registerPass([&] { return dxil::ShaderFlagsAnalysis(); });
   });
 }
 
 bool DirectXTargetMachine::addPassesToEmitFile(
     PassManagerBase &PM, raw_pwrite_stream &Out, raw_pwrite_stream *DwoOut,
     CodeGenFileType FileType, bool DisableVerify,
     MachineModuleInfoWrapperPass *MMIWP) {

llvm/test/CodeGen/DirectX/UAVMetadata.ll

 ; RUN: opt -S -dxil-metadata-emit < %s | FileCheck %s
 ; RUN: opt -S --passes="print-dxil-resource" < %s 2>&1 | FileCheck %s --check-prefix=PRINT
+; RUN: opt -S -dxil-metadata-emit < %s | opt -S --passes="print-dxil-resource-from-dxil-metadata" 2>&1 | FileCheck %s --check-prefix=PRINT
+
 ; RUN: llc %s --filetype=asm -o - < %s 2>&1 | FileCheck %s --check-prefixes=CHECK,PRINT
 
 target datalayout = "e-m:e-p:32:32-i1:32-i8:8-i16:16-i32:32-i64:64-f16:16-f32:32-f64:64-n8:16:32:64"
 target triple = "dxil-pc-shadermodel6.0-library"
 
 %"class.hlsl::RWBuffer" = type { ptr }
 
 

llvm/test/CodeGen/DirectX/cbuf.ll

 ; RUN: opt -S -dxil-metadata-emit < %s | FileCheck %s --check-prefix=DXILMD
 ; RUN: opt -S --passes="print-dxil-resource" < %s 2>&1 | FileCheck %s --check-prefix=PRINT
+; RUN: opt -S -dxil-metadata-emit < %s | opt -S --passes="print-dxil-resource-from-dxil-metadata" 2>&1 | FileCheck %s --check-prefix=PRINT
 
 target datalayout = "e-m:e-p:32:32-i1:32-i8:8-i16:16-i32:32-i64:64-f16:16-f32:32-f64:64-n8:16:32:64"
 target triple = "dxil-unknown-shadermodel6.7-library"
 
 ; Make sure the size is 24 = 16 + 8 (float,i32,double -> 16 and int2 -> 8)
 ; DXILMD:!{i32 0, ptr @A.cb., !"", i32 1, i32 2, i32 1, i32 24}
 
 ; Make sure match register(b2, space1) with ID 0.