[HLSL] CodeGen HLSL Resource annotations

[beanz]

Summary

HLSL Resource types need special annotations that the backend will use
to build out metadata and resource annotations that are required by
DirectX and Vulkan drivers in order to provide correct data bindings
for shader exeuction.

This patch adds some of the required data for unordered-access-views
(UAV) resource binding into the module flags. This data will evolve
over time to cover all the required use cases, but this should get
things started.

Depends on D130018.

[beanz]

Updating based on feedback from @python3kgae provided via chat. Because of the odd nature of DXIL, the DirectX backend needs to filter the module flags, for that reason it is easier to use a named metadata entry which can just be deleted during codegen. This should be fine and have no real impact on other targets.

[aaron.ballman]

clang/lib/CodeGen/CGHLSLRuntime.cpp

59

62–66

78

Add an assertion that ResourceMD is always nonnull by this point?

[beanz]

Updates based on PR feedback. Thanks @aaron.ballman!

[aaron.ballman]

LGTM!

clang/lib/CodeGen/CGDeclCXX.cpp

 //===--- CGDeclCXX.cpp - Emit LLVM Code for C++ declarations --------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This contains code dealing with code generation of C++ declarations
 //
 //===----------------------------------------------------------------------===//
 
 #include "CGCXXABI.h"
+#include "CGHLSLRuntime.h"
 #include "CGObjCRuntime.h"
 #include "CGOpenMPRuntime.h"
 #include "CodeGenFunction.h"
 #include "TargetInfo.h"
 #include "clang/AST/Attr.h"
 #include "clang/Basic/LangOptions.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/IR/Intrinsics.h"
void CodeGenFunction::GenerateCXXGlobalVarDeclInitFunc(llvm::Function *Fn,
   if (Addr->hasWeakLinkage() || Addr->hasLinkOnceLinkage() ||
       (D->getTLSKind() == VarDecl::TLS_Dynamic &&
        isTemplateInstantiation(D->getTemplateSpecializationKind()))) {
     EmitCXXGuardedInit(*D, Addr, PerformInit);
   } else {
     EmitCXXGlobalVarDeclInit(*D, Addr, PerformInit);
   }
 
+  if (getLangOpts().HLSL)
+    CGM.getHLSLRuntime().annotateHLSLResource(D, Addr);
+
   FinishFunction();
 }
 
 void
 CodeGenFunction::GenerateCXXGlobalInitFunc(llvm::Function *Fn,
                                            ArrayRef<llvm::Function *> Decls,
                                            ConstantAddress Guard) {
   {

clang/lib/CodeGen/CGHLSLRuntime.h

 // subclasses of this implement code generation for specific HLSL
 // runtime libraries.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_LIB_CODEGEN_CGHLSLRUNTIME_H
 #define LLVM_CLANG_LIB_CODEGEN_CGHLSLRUNTIME_H
 
+#include "clang/Basic/HLSLRuntime.h"
+
+namespace llvm {
+class Value;
+class GlobalVariable;
+} // namespace llvm
 namespace clang {
+class CallExpr;
+class Type;
+class VarDecl;
 
 namespace CodeGen {
 
 class CodeGenModule;
 
 class CGHLSLRuntime {
 protected:
   CodeGenModule &CGM;
+  uint32_t ResourceCounters[static_cast<uint32_t>(
+      hlsl::ResourceClass::NumClasses)] = {0};
 
 public:
   CGHLSLRuntime(CodeGenModule &CGM) : CGM(CGM) {}
   virtual ~CGHLSLRuntime() {}
 
+  void annotateHLSLResource(const VarDecl *D, llvm::GlobalVariable *GV);
+
   void finishCodeGen();
 };
 
 } // namespace CodeGen
 } // namespace clang
 
 #endif

clang/lib/CodeGen/CGHLSLRuntime.cpp

 #include "CGHLSLRuntime.h"
 #include "CodeGenModule.h"
 #include "clang/Basic/TargetOptions.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 
 using namespace clang;
 using namespace CodeGen;
+using namespace hlsl;
 using namespace llvm;
 
 namespace {
 void addDxilValVersion(StringRef ValVersionStr, llvm::Module &M) {
   // The validation of ValVersionStr is done at HLSLToolChain::TranslateArgs.
   // Assume ValVersionStr is legal here.
   VersionTuple Version;
   if (Version.tryParse(ValVersionStr) || Version.getBuild() ||

 } // namespace
 
 void CGHLSLRuntime::finishCodeGen() {
   auto &TargetOpts = CGM.getTarget().getTargetOpts();
 
   llvm::Module &M = CGM.getModule();
   addDxilValVersion(TargetOpts.DxilValidatorVersion, M);
 }
+
+void CGHLSLRuntime::annotateHLSLResource(const VarDecl *D, GlobalVariable *GV) {
+  const Type *Ty = D->getType()->getPointeeOrArrayElementType();
+  if (!Ty)
+    return;
+  const auto *RD = Ty->getAsCXXRecordDecl();
+  if (!RD)
+    return;
+  const auto *Attr = RD->getAttr<HLSLResourceAttr>();
+  if (!Attr)
+    return;
+
+  HLSLResourceAttr::ResourceClass RC = Attr->getResourceType();
+  uint32_t Counter = ResourceCounters[static_cast<uint32_t>(RC)]++;
+
+  NamedMDNode *ResourceMD = nullptr;
+  switch (RC) {
+  case HLSLResourceAttr::ResourceClass::UAV:
+    ResourceMD = CGM.getModule().getOrInsertNamedMetadata("hlsl.uavs");
+    break;
+  default:
+    assert(false && "Unsupported buffer type!");
+    return;
+  }
+
+  assert(ResourceMD != nullptr &&
+         "ResourceMD must have been set by the switch above.");
+
+  auto &Ctx = CGM.getModule().getContext();
+  IRBuilder<> B(Ctx);
+  QualType QT(Ty, 0);
+  ResourceMD->addOperand(MDNode::get(
+      Ctx, {ValueAsMetadata::get(GV), MDString::get(Ctx, QT.getAsString()),
+            ConstantAsMetadata::get(B.getInt32(Counter))}));
+}

clang/test/CodeGenHLSL/builtins/RWBuffer-annotations.hlsl

+// RUN: %clang_cc1 -triple dxil-pc-shadermodel6.0-compute -x hlsl -emit-llvm -disable-llvm-passes -o - %s | FileCheck %s 
+
+RWBuffer<float> Buffer1;
+RWBuffer<vector<float, 4> > BufferArray[4];
+
+[numthreads(1,1,1)]
+void main() {
+}
+
+// CHECK: !hlsl.uavs = !{![[Single:[0-9]+]], ![[Array:[0-9]+]]}
+// CHECK-DAG: ![[Single]] = !{ptr @"?Buffer1@@3V?$RWBuffer@M@hlsl@@A", !"RWBuffer<float>", i32 0}
+// CHECK-DAG: ![[Array]] = !{ptr @"?BufferArray@@3PAV?$RWBuffer@T?$__vector@M$03@__clang@@@hlsl@@A", !"RWBuffer<vector<float, 4> >", i32 1}