[llvm] [SPIRV][HLSL] Implement CBuffer access lowering pass (PR #159136)

[Nathan Gauër via llvm-commits]

================
@@ -0,0 +1,172 @@
+//===- SPIRVCBufferAccess.cpp - Translate CBuffer Loads
+//--------------------===//
+//
+// 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 pass replaces all accesses to constant buffer global variables with
+// accesses to the proper SPIR-V resource. It's designed to run after the
+// DXIL preparation passes and before the main SPIR-V legalization passes.
+//
+// The pass operates as follows:
+// 1. It finds all constant buffers by looking for the `!hlsl.cbs` metadata.
+// 2. For each cbuffer, it finds the global variable holding the resource handle
+//    and the global variables for each of the cbuffer's members.
+// 3. For each member variable, it creates a call to the
+//    `llvm.spv.resource.getpointer` intrinsic. This intrinsic takes the
+//    resource handle and the member's index within the cbuffer as arguments.
+//    The result is a pointer to that member within the SPIR-V resource.
+// 4. It then replaces all uses of the original member global variable with the
+//    pointer returned by the `getpointer` intrinsic. This effectively retargets
+//    all loads and GEPs to the new resource pointer.
+// 5. Finally, it cleans up by deleting the original global variables and the
+//    `!hlsl.cbs` metadata.
+//
+// This approach allows subsequent passes, like SPIRVEmitIntrinsics, to
+// correctly handle GEPs that operate on the result of the `getpointer` call,
+// folding them into a single OpAccessChain instruction.
+//
+//===----------------------------------------------------------------------===//
+
+#include "SPIRVCBufferAccess.h"
+#include "SPIRV.h"
+#include "llvm/Frontend/HLSL/CBuffer.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/IntrinsicsSPIRV.h"
+#include "llvm/IR/Module.h"
+
+#define DEBUG_TYPE "spirv-cbuffer-access"
+using namespace llvm;
+
+// Finds the single instruction that defines the resource handle. This is
+// typically a call to `llvm.spv.resource.handlefrombinding`.
+static Instruction *findHandleDef(GlobalVariable *HandleVar) {
+  for (User *U : HandleVar->users()) {
+    if (auto *SI = dyn_cast<StoreInst>(U)) {
+      if (auto *I = dyn_cast<Instruction>(SI->getValueOperand())) {
+        return I;
+      }
+    }
+  }
+  return nullptr;
+}
+
+static bool replaceCBufferAccesses(Module &M) {
+  std::optional<hlsl::CBufferMetadata> CBufMD = hlsl::CBufferMetadata::get(M);
+  if (!CBufMD)
+    return false;
+
+  for (const hlsl::CBufferMapping &Mapping : *CBufMD) {
+    Instruction *HandleDef = findHandleDef(Mapping.Handle);
+    if (!HandleDef) {
+      // If there's no handle definition, it might be because the cbuffer is
+      // unused. In this case, we can just clean up the globals.
+      if (Mapping.Handle->use_empty()) {
+        for (const auto &Member : Mapping.Members) {
+          if (Member.GV->use_empty()) {
+            Member.GV->eraseFromParent();
+          }
+        }
+        Mapping.Handle->eraseFromParent();
+      }
+      continue;
+    }
+
+    // The handle definition should dominate all uses of the cbuffer members.
+    // We'll insert our getpointer calls right after it.
+    IRBuilder<> Builder(HandleDef->getNextNode());
+
+    for (uint32_t Index = 0; Index < Mapping.Members.size(); ++Index) {
+      GlobalVariable *MemberGV = Mapping.Members[Index].GV;
+      if (MemberGV->use_empty()) {
+        continue;
+      }
+
+      // Create the getpointer intrinsic call.
+      Value *IndexVal = Builder.getInt32(Index);
+      Type *PtrType = MemberGV->getType();
+      Value *GetPointerCall = Builder.CreateIntrinsic(
+          PtrType, Intrinsic::spv_resource_getpointer, {HandleDef, IndexVal});
+
+      // We cannot use replaceAllUsesWith here because some uses may be
+      // ConstantExprs, which cannot be replaced with non-constants.
+      SmallVector<User *, 4> Users(MemberGV->users());
+      for (User *U : Users) {
+        if (auto *CE = dyn_cast<ConstantExpr>(U)) {
+          SmallVector<Instruction *, 4> Insts;
+          std::function<void(ConstantExpr *)> findInstructions =
+              [&](ConstantExpr *Const) {
+                for (User *ConstU : Const->users()) {
+                  if (auto *ConstCE = dyn_cast<ConstantExpr>(ConstU)) {
+                    findInstructions(ConstCE);
+                  } else if (auto *I = dyn_cast<Instruction>(ConstU)) {
+                    Insts.push_back(I);
+                  }
+                }
+              };
+          findInstructions(CE);
+
+          for (Instruction *I : Insts) {
+            Instruction *NewInst = CE->getAsInstruction();
+            NewInst->insertBefore(I->getIterator());
+            I->replaceUsesOfWith(CE, NewInst);
+            NewInst->replaceUsesOfWith(MemberGV, GetPointerCall);
+          }
+        } else {
+          U->replaceUsesOfWith(MemberGV, GetPointerCall);
+        }
+      }
+    }
+  }
+
+  // Now that all uses are replaced, clean up the globals and metadata.
+  for (const hlsl::CBufferMapping &Mapping : *CBufMD) {
+    for (const auto &Member : Mapping.Members) {
+      Member.GV->eraseFromParent();
+    }
+    // Erase the stores to the handle variable before erasing the handle itself.
+    SmallVector<Instruction *, 4> HandleStores;
+    for (User *U : Mapping.Handle->users()) {
+      if (auto *SI = dyn_cast<StoreInst>(U)) {
+        HandleStores.push_back(SI);
+      }
+    }
+    for (Instruction *I : HandleStores) {
+      I->eraseFromParent();
+    }
+    Mapping.Handle->eraseFromParent();
+  }
+
+  CBufMD->eraseFromModule();
+  return true;
+}
+
+PreservedAnalyses SPIRVCBufferAccess::run(Module &M,
+                                          ModuleAnalysisManager &AM) {
+  if (replaceCBufferAccesses(M)) {
+    return PreservedAnalyses::none();
+  }
+  return PreservedAnalyses::all();
+}
+
+namespace {
----------------
Keenuts wrote:

nit:
you could put everything in the anonymous namespace, and since you got code + function in the cpp, sink the `::run` code into the class definition to keep it local no?

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