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

[Steven Perron 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 {
----------------
s-perron wrote:

I'm not sure what you mean. I want the pass for the new pass manager and one for the legacy pass manager. I have on static function that both can share. Only the class is in the anonymous namespace because it seems like llvm prefers using static function instead of the anonymous namespace. 

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