[DirectX] Generate `dx.resources` metadata entry

[beanz]

Summary

This code adds initial support for generating the HLSL resources
metadata entries. It has a lot of FIXMEs laying around because there
is a lot more work to do here, but this lays a solid groundwork and can
accurately handle some trivial cases.

I've filed a swath of issues covering the deficiencies here and left the
issues in comments so that we can easily follow them.

One big change to make sooner rather than later is to move some of this
code into a new libLLVMFrontendHLSL so that we can share it with the
Clang CodeGen layer.

[python3kgae]

llvm/lib/Target/DirectX/DXILResource.h

143–145

Maybe another class like

class HLSL_UAVs {
public:

llvm::SmallVector<UAVResource> UAVs;
HLSL_UAVs() {... call collectUAVs }
~HLSL_UAVs() { remove the "hlsl.uavs" metadata}

private:

void collectUAVs();

};

[beanz]

llvm/lib/Target/DirectX/DXILResource.h

143–145

Given that the current class only handles UAVs, I'm not sure the benefit of an extra class. That might make sense as a refactoring as we support more resource types, but for now it seems like an unnecessary abstraction.

[python3kgae]

llvm/lib/Target/DirectX/DXILResource.cpp

123–130

This could be in ResourceBase.
Refactoring when we got other resource classes too?

[beanz]

Pulling writing base entries out into ResourceBase::write based on feedback from @python3kgae

[python3kgae]

UAVResource::parseSourceType is really ugly.
It would be nice to fix it in the next PR.

Also a slight format issue.

llvm/lib/Target/DirectX/DXILResource.cpp

147

llvm/lib/Target/DirectX/CMakeLists.txt

add_llvm_target(DirectXCodeGen
   DirectXInstrInfo.cpp
   DirectXRegisterInfo.cpp
   DirectXSubtarget.cpp
   DirectXTargetMachine.cpp
   DXILMetadata.cpp
   DXILOpBuilder.cpp
   DXILOpLowering.cpp
   DXILPrepare.cpp
+  DXILResource.cpp
   DXILTranslateMetadata.cpp
   PointerTypeAnalysis.cpp
 
   LINK_COMPONENTS
   AsmPrinter
   Core
   SelectionDAG
   Support

llvm/lib/Target/DirectX/DXILResource.h

+//===- DXILResource.h - DXIL Resource 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 Resources.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TARGET_DIRECTX_DXILRESOURCE_H
+#define LLVM_TARGET_DIRECTX_DXILRESOURCE_H
+
+#include "llvm/ADT/Optional.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/IR/Metadata.h"
+#include <cstdint>
+
+namespace llvm {
+class Module;
+class GlobalVariable;
+
+namespace dxil {
+
+// FIXME: Ultimately this class and some of these utilities should be moved into
+// a new LLVMFrontendHLSL library so that they can be reused in Clang.
+// See issue https://github.com/llvm/llvm-project/issues/58000.
+class FrontendResource {
+  MDNode *Entry;
+
+public:
+  FrontendResource(MDNode *E) : Entry(E) {
+    assert(Entry->getNumOperands() == 3 && "Unexpected metadata shape");
+  }
+
+  GlobalVariable *getGlobalVariable();
+  StringRef getSourceType();
+  Constant *getID();
+};
+
+class ResourceBase {
+protected:
+  uint32_t ID;
+  GlobalVariable *GV;
+  StringRef Name;
+  uint32_t Space;
+  uint32_t LowerBound;
+  uint32_t RangeSize;
+  ResourceBase(uint32_t I, FrontendResource R);
+
+  void write(LLVMContext &Ctx, MutableArrayRef<Metadata *> Entries);
+
+  // The value ordering of this enumeration is part of the DXIL ABI. Elements
+  // can only be added to the end, and not removed.
+  enum class Kinds : uint32_t {
+    Invalid = 0,
+    Texture1D,
+    Texture2D,
+    Texture2DMS,
+    Texture3D,
+    TextureCube,
+    Texture1DArray,
+    Texture2DArray,
+    Texture2DMSArray,
+    TextureCubeArray,
+    TypedBuffer,
+    RawBuffer,
+    StructuredBuffer,
+    CBuffer,
+    Sampler,
+    TBuffer,
+    RTAccelerationStructure,
+    FeedbackTexture2D,
+    FeedbackTexture2DArray,
+    NumEntries,
+  };
+
+  // The value ordering of this enumeration is part of the DXIL ABI. Elements
+  // can only be added to the end, and not removed.
+  enum class ComponentType : uint32_t {
+    Invalid = 0,
+    I1,
+    I16,
+    U16,
+    I32,
+    U32,
+    I64,
+    U64,
+    F16,
+    F32,
+    F64,
+    SNormF16,
+    UNormF16,
+    SNormF32,
+    UNormF32,
+    SNormF64,
+    UNormF64,
+    PackedS8x32,
+    PackedU8x32,
+    LastEntry
+  };
+
+public:
+  struct ExtendedProperties {
+    llvm::Optional<ComponentType> ElementType;
+
+    // The value ordering of this enumeration is part of the DXIL ABI. Elements
+    // can only be added to the end, and not removed.
+    enum Tags : uint32_t {
+      TypedBufferElementType = 0,
+      StructuredBufferElementStride,
+      SamplerFeedbackKind,
+      Atomic64Use
+    };
+
+    MDNode *write(LLVMContext &Ctx);
+  };
+};
+
+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, FrontendResource R);
+
+  MDNode *write();
+};
+
+// 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 {
+  Module &Mod;
+  llvm::SmallVector<UAVResource> UAVs;
+
+  void collectUAVs();
+
+public:
+  Resources(Module &M) : Mod(M) { collectUAVs(); }
+
+  void write();
+};
+
+} // namespace dxil
+} // namespace llvm
+
+#endif // LLVM_TARGET_DIRECTX_DXILRESOURCE_H

llvm/lib/Target/DirectX/DXILResource.cpp

+//===- DXILResource.cpp - DXIL Resource 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 Resources.
+///
+//===----------------------------------------------------------------------===//
+
+#include "DXILResource.h"
+#include "llvm/ADT/StringSwitch.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Metadata.h"
+#include "llvm/IR/Module.h"
+
+using namespace llvm;
+using namespace llvm::dxil;
+
+GlobalVariable *FrontendResource::getGlobalVariable() {
+  return cast<GlobalVariable>(
+      cast<ConstantAsMetadata>(Entry->getOperand(0))->getValue());
+}
+
+StringRef FrontendResource::getSourceType() {
+  return cast<MDString>(Entry->getOperand(1))->getString();
+}
+
+Constant *FrontendResource::getID() {
+  return cast<ConstantAsMetadata>(Entry->getOperand(2))->getValue();
+}
+
+void Resources::collectUAVs() {
+  NamedMDNode *Entry = Mod.getNamedMetadata("hlsl.uavs");
+  if (!Entry || Entry->getNumOperands() == 0)
+    return;
+
+  uint32_t Counter = 0;
+  for (auto *UAV : Entry->operands()) {
+    UAVs.push_back(UAVResource(Counter++, FrontendResource(cast<MDNode>(UAV))));
+  }
+}
+
+ResourceBase::ResourceBase(uint32_t I, FrontendResource R)
+    : ID(I), GV(R.getGlobalVariable()), Name(""), Space(0), LowerBound(0),
+      RangeSize(1) {
+  if (auto *ArrTy = dyn_cast<ArrayType>(GV->getInitializer()->getType()))
+    RangeSize = ArrTy->getNumElements();
+}
+
+UAVResource::UAVResource(uint32_t I, FrontendResource R)
+    : ResourceBase(I, R), Shape(Kinds::Invalid), GloballyCoherent(false),
+      HasCounter(false), IsROV(false), ExtProps() {
+  parseSourceType(R.getSourceType());
+}
+
+// FIXME: Capture this in HLSL source. I would go do this right now, but I want
+// to get this in first so that I can make sure to capture all the extra
+// information we need to remove the source type string from here (See issue:
+// https://github.com/llvm/llvm-project/issues/57991).
+void UAVResource::parseSourceType(StringRef S) {
+  IsROV = S.startswith("RasterizerOrdered");
+  if (IsROV)
+    S = S.substr(strlen("RasterizerOrdered"));
+  if (S.startswith("RW"))
+    S = S.substr(strlen("RW"));
+
+  // Note: I'm deliberately not handling any of the Texture buffer types at the
+  // moment. I want to resolve the issue above before adding Texture or Sampler
+  // support.
+  Shape = StringSwitch<ResourceBase::Kinds>(S)
+              .StartsWith("Buffer<", Kinds::TypedBuffer)
+              .StartsWith("ByteAddressBuffer<", Kinds::RawBuffer)
+              .StartsWith("StructuredBuffer<", Kinds::StructuredBuffer)
+              .Default(Kinds::Invalid);
+  assert(Shape != Kinds::Invalid && "Unsupported buffer type");
+
+  S = S.substr(S.find("<") + 1);
+
+  constexpr size_t PrefixLen = StringRef("vector<").size();
+  if (S.startswith("vector<"))
+    S = S.substr(PrefixLen, S.find(",") - PrefixLen);
+  else
+    S = S.substr(0, S.find(">"));
+
+  ComponentType ElTy = StringSwitch<ResourceBase::ComponentType>(S)
+                           .Case("bool", ComponentType::I1)
+                           .Case("int16_t", ComponentType::I16)
+                           .Case("uint16_t", ComponentType::U16)
+                           .Case("int32_t", ComponentType::I32)
+                           .Case("uint32_t", ComponentType::U32)
+                           .Case("int64_t", ComponentType::I64)
+                           .Case("uint64_t", ComponentType::U64)
+                           .Case("half", ComponentType::F16)
+                           .Case("float", ComponentType::F32)
+                           .Case("double", ComponentType::F64)
+                           .Default(ComponentType::Invalid);
+  if (ElTy != ComponentType::Invalid)
+    ExtProps.ElementType = ElTy;
+}
+
+MDNode *ResourceBase::ExtendedProperties::write(LLVMContext &Ctx) {
+  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))));
+  }
+  if (Entries.empty())
+    return nullptr;
+  return MDNode::get(Ctx, Entries);
+}
+
+void ResourceBase::write(LLVMContext &Ctx,
+                         MutableArrayRef<Metadata *> Entries) {
+  IRBuilder<> B(Ctx);
+  Entries[0] = ConstantAsMetadata::get(B.getInt32(ID));
+  Entries[1] = ConstantAsMetadata::get(GV);
+  Entries[2] = MDString::get(Ctx, Name);
+  Entries[3] = ConstantAsMetadata::get(B.getInt32(Space));
+  Entries[4] = ConstantAsMetadata::get(B.getInt32(LowerBound));
+  Entries[5] = ConstantAsMetadata::get(B.getInt32(RangeSize));
+}
+
+MDNode *UAVResource::write() {
+  auto &Ctx = GV->getContext();
+  IRBuilder<> B(Ctx);
+  Metadata *Entries[11];
+  ResourceBase::write(Ctx, Entries);
+  Entries[6] =
+      ConstantAsMetadata::get(B.getInt32(static_cast<uint32_t>(Shape)));
+  Entries[7] = ConstantAsMetadata::get(B.getInt1(GloballyCoherent));
+  Entries[8] = ConstantAsMetadata::get(B.getInt1(HasCounter));
+  Entries[9] = ConstantAsMetadata::get(B.getInt1(IsROV));
+  Entries[10] = ExtProps.write(Ctx);
+  return MDNode::get(Ctx, Entries);
+}
+
+void Resources::write() {
+  Metadata *ResourceMDs[4] = {nullptr, nullptr, nullptr, nullptr};
+  SmallVector<Metadata *> UAVMDs;
+  for (auto &UAV : UAVs)
+    UAVMDs.emplace_back(UAV.write());
+
+  if (!UAVMDs.empty())
+    ResourceMDs[1] = MDNode::get(Mod.getContext(), UAVMDs);
+
+  NamedMDNode *DXResMD = Mod.getOrInsertNamedMetadata("dx.resources");
+  DXResMD->addOperand(MDNode::get(Mod.getContext(), ResourceMDs));
+
+  NamedMDNode *Entry = Mod.getNamedMetadata("hlsl.uavs");
+  if (Entry)
+    Entry->eraseFromParent();
+}

llvm/lib/Target/DirectX/DXILTranslateMetadata.cpp

 //===- DXILTranslateMetadata.cpp - Pass to emit DXIL metadata ---*- 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
 //
 //===----------------------------------------------------------------------===//
 ///
 //===----------------------------------------------------------------------===//
 
 #include "DXILMetadata.h"
+#include "DXILResource.h"
 #include "DirectX.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Pass.h"
 

 } // namespace
 
 bool DXILTranslateMetadata::runOnModule(Module &M) {
 
   dxil::ValidatorVersionMD ValVerMD(M);
   if (ValVerMD.isEmpty())
     ValVerMD.update(VersionTuple(1, 0));
   dxil::createShaderModelMD(M);
+
+  dxil::Resources Res(M);
+  Res.write();
   return false;
 }
 
 char DXILTranslateMetadata::ID = 0;
 
 ModulePass *llvm::createDXILTranslateMetadataPass() {
   return new DXILTranslateMetadata();
 }
 
 INITIALIZE_PASS(DXILTranslateMetadata, "dxil-metadata-emit",
                 "DXIL Metadata Emit", false, false)

llvm/test/CodeGen/DirectX/UAVMetadata.ll

+; RUN: opt -S -dxil-metadata-emit < %s | FileCheck %s
+; ModuleID = '/home/cbieneman/dev/shuffle.hlsl'
+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-compute"
+
+%"class.hlsl::RWBuffer" = type { ptr }
+
+@Zero = local_unnamed_addr global %"class.hlsl::RWBuffer" zeroinitializer, align 4
+@One = local_unnamed_addr global %"class.hlsl::RWBuffer" zeroinitializer, align 4
+@Two = local_unnamed_addr global %"class.hlsl::RWBuffer" zeroinitializer, align 4
+@Three = local_unnamed_addr global [2 x %"class.hlsl::RWBuffer"] zeroinitializer, align 4
+@Four = local_unnamed_addr global %"class.hlsl::RWBuffer" zeroinitializer, align 4
+@Five = local_unnamed_addr global %"class.hlsl::RWBuffer" zeroinitializer, align 4
+@Six = local_unnamed_addr global %"class.hlsl::RWBuffer" zeroinitializer, align 4
+@Seven = local_unnamed_addr global %"class.hlsl::RWBuffer" zeroinitializer, align 4
+@Eight = local_unnamed_addr global %"class.hlsl::RWBuffer" zeroinitializer, align 4
+@Nine = local_unnamed_addr global %"class.hlsl::RWBuffer" zeroinitializer, align 4
+
+
+!hlsl.uavs = !{!0, !1, !2, !3, !4, !5, !6, !7, !8, !9}
+
+!0 = !{ptr @Zero, !"RWBuffer<half>", i32 0}
+!1 = !{ptr @One, !"Buffer<vector<float,4>>", i32 1}
+!2 = !{ptr @Two, !"Buffer<double>", i32 2}
+!3 = !{ptr @Three, !"Buffer<bool>", i32 3}
+!4 = !{ptr @Four, !"ByteAddressBuffer<int16_t>", i32 4}
+!5 = !{ptr @Five, !"StructuredBuffer<uint16_t>", i32 5}
+!6 = !{ptr @Six, !"RasterizerOrderedBuffer<int32_t>", i32 6}
+!7 = !{ptr @Seven, !"RasterizerOrderedStructuredBuffer<uint32_t>", i32 7}
+!8 = !{ptr @Eight, !"RasterizerOrderedByteAddressBuffer<int64_t>", i32 8}
+!9 = !{ptr @Nine, !"RWBuffer<uint64_t>", i32 9}
+
+
+; CHECK: !dx.resources = !{[[ResList:[!][0-9]+]]}
+
+; CHECK: [[ResList]] = !{null, [[UAVList:[!][0-9]+]], null, null}
+; CHECK: [[UAVList]] = !{[[Zero:[!][0-9]+]], [[One:[!][0-9]+]],
+; CHECK-SAME: [[Two:[!][0-9]+]], [[Three:[!][0-9]+]], [[Four:[!][0-9]+]], 
+; CHECK-SAME: [[Five:[!][0-9]+]], [[Six:[!][0-9]+]], [[Seven:[!][0-9]+]],
+; CHECK-SAME: [[Eight:[!][0-9]+]], [[Nine:[!][0-9]+]]}
+; CHECK: [[Zero]] = !{i32 0, ptr @Zero, !"", i32 0, i32 0, i32 1, i32 10, i1 false, i1 false, i1 false, [[Half:[!][0-9]+]]}
+; CHECK: [[Half]] = !{i32 0, i32 8}
+; CHECK: [[One]] = !{i32 1, ptr @One, !"", i32 0, i32 0, i32 1, i32 10, i1 false, i1 false, i1 false, [[Float:[!][0-9]+]]}
+; CHECK: [[Float]] = !{i32 0, i32 9}
+; CHECK: [[Two]] = !{i32 2, ptr @Two, !"", i32 0, i32 0, i32 1, i32 10, i1 false, i1 false, i1 false, [[Double:[!][0-9]+]]}
+; CHECK: [[Double]] = !{i32 0, i32 10}
+; CHECK: [[Three]] = !{i32 3, ptr @Three, !"", i32 0, i32 0, i32 2, i32 10, i1 false, i1 false, i1 false, [[Bool:[!][0-9]+]]}
+; CHECK: [[Bool]] = !{i32 0, i32 1}
+; CHECK: [[Four]] = !{i32 4, ptr @Four, !"", i32 0, i32 0, i32 1, i32 11, i1 false, i1 false, i1 false, [[I16:[!][0-9]+]]}
+; CHECK: [[I16]] = !{i32 0, i32 2}
+; CHECK: [[Five]] = !{i32 5, ptr @Five, !"", i32 0, i32 0, i32 1, i32 12, i1 false, i1 false, i1 false, [[U16:[!][0-9]+]]}
+; CHECK: [[U16]] = !{i32 0, i32 3}
+; CHECK: [[Six]] = !{i32 6, ptr @Six, !"", i32 0, i32 0, i32 1, i32 10, i1 false, i1 false, i1 true, [[I32:[!][0-9]+]]}
+; CHECK: [[I32]] = !{i32 0, i32 4}
+; CHECK: [[Seven]] = !{i32 7, ptr @Seven, !"", i32 0, i32 0, i32 1, i32 12, i1 false, i1 false, i1 true, [[U32:[!][0-9]+]]}
+; CHECK: [[U32]] = !{i32 0, i32 5}
+; CHECK: [[Eight]] = !{i32 8, ptr @Eight, !"", i32 0, i32 0, i32 1, i32 11, i1 false, i1 false, i1 true, [[I64:[!][0-9]+]]}
+; CHECK: [[I64]] = !{i32 0, i32 6}
+; CHECK: [[Nine]] = !{i32 9, ptr @Nine, !"", i32 0, i32 0, i32 1, i32 10, i1 false, i1 false, i1 false, [[U64:[!][0-9]+]]}
+; CHECK: [[U64]] = !{i32 0, i32 7}