[llvm] [DirectX] Implement DXILResourceImplicitBinding pass (PR #138043)

[Ashley Coleman via llvm-commits]

================
@@ -0,0 +1,181 @@
+//===- DXILResourceImplicitBinding.cpp -----------------------------------===//
+//
+// 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 "DXILResourceImplicitBinding.h"
+#include "DirectX.h"
+#include "llvm/ADT/APInt.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Analysis/DXILResource.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/DiagnosticInfo.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/IntrinsicsDirectX.h"
+#include "llvm/IR/Module.h"
+#include "llvm/InitializePasses.h"
+#include <cstdint>
+
+#define DEBUG_TYPE "dxil-resource-implicit-binding"
+
+using namespace llvm;
+using namespace llvm::dxil;
+
+namespace {
+
+static void diagnoseImplicitBindingNotFound(CallInst *ImplBindingCall) {
+  Function *F = ImplBindingCall->getFunction();
+  LLVMContext &Context = F->getParent()->getContext();
+  // FIXME: include the name of the resource in the error message
+  // (llvm/llvm-project#137868)
+  Context.diagnose(
+      DiagnosticInfoGenericWithLoc("resource cannot be allocated", *F,
+                                   ImplBindingCall->getDebugLoc(), DS_Error));
+}
+
+static bool assignBindings(Module &M, DXILResourceBindingInfo &DRBI,
+                           DXILResourceTypeMap &DRTM) {
+  struct ImplicitBindingCall {
+    int OrderID;
+    CallInst *Call;
+    ImplicitBindingCall(int OrderID, CallInst *Call)
+        : OrderID(OrderID), Call(Call) {}
+  };
+  SmallVector<ImplicitBindingCall> Calls;
+  SmallVector<Function *> FunctionsToMaybeRemove;
+
+  // collect all of the llvm.dx.resource.handlefromImplicitbinding calls
+  for (Function &F : M.functions()) {
+    if (!F.isDeclaration())
+      continue;
+
+    if (F.getIntrinsicID() != Intrinsic::dx_resource_handlefromimplicitbinding)
+      continue;
+
+    for (User *U : F.users()) {
+      if (CallInst *CI = dyn_cast<CallInst>(U)) {
+        int OrderID = cast<ConstantInt>(CI->getArgOperand(0))->getZExtValue();
+        Calls.emplace_back(OrderID, CI);
+      }
+    }
+    FunctionsToMaybeRemove.emplace_back(&F);
+  }
+
+  // sort all the collected implicit bindings by OrderID
+  llvm::stable_sort(
+      Calls, [](auto &LHS, auto &RHS) { return LHS.OrderID < RHS.OrderID; });
+
+  // iterate over sorted calls, find binding for each new OrderID and replace
+  // each call with dx_resource_handlefrombinding using the new binding
+  int LastOrderID = -1;
+  llvm::TargetExtType *HandleTy = nullptr;
+  ConstantInt *RegSlotOp = nullptr;
+  bool AllBindingsAssigned = true;
+  bool Changed = false;
+
+  for (ImplicitBindingCall &IB : Calls) {
+    IRBuilder<> Builder(IB.Call);
+
+    if (IB.OrderID != LastOrderID) {
+      LastOrderID = IB.OrderID;
+      HandleTy = cast<TargetExtType>(IB.Call->getType());
+      ResourceTypeInfo &RTI = DRTM[HandleTy];
+
+      uint32_t Space =
+          cast<ConstantInt>(IB.Call->getArgOperand(1))->getZExtValue();
+      int32_t Size =
+          cast<ConstantInt>(IB.Call->getArgOperand(2))->getZExtValue();
+
+      std::optional<uint32_t> RegSlot =
+          DRBI.findAvailableBinding(RTI.getResourceClass(), Space, Size);
+      if (!RegSlot) {
+        diagnoseImplicitBindingNotFound(IB.Call);
+        AllBindingsAssigned = false;
+        continue;
+      }
+      RegSlotOp = ConstantInt::get(Builder.getInt32Ty(), RegSlot.value());
+    }
+
+    if (!RegSlotOp)
+      continue;
+
+    auto *NewCall = Builder.CreateIntrinsic(
+        HandleTy, Intrinsic::dx_resource_handlefrombinding,
+        {IB.Call->getOperand(1),   /* space */
+         RegSlotOp,                /* register slot */
+         IB.Call->getOperand(2),   /* size */
+         IB.Call->getOperand(3),   /* index */
+         IB.Call->getOperand(4)}); /* non-uniform flag */
+    IB.Call->replaceAllUsesWith(NewCall);
+    IB.Call->eraseFromParent();
+    Changed = true;
+  }
+
+  for (Function *F : FunctionsToMaybeRemove) {
+    if (F->user_empty()) {
+      F->eraseFromParent();
+      Changed = true;
+    }
+  }
+
+  DRBI.setHasImplicitBinding(!AllBindingsAssigned);
+  return Changed;
+}
+
+} // end anonymous namespace
+
+PreservedAnalyses DXILResourceImplicitBinding::run(Module &M,
+                                                   ModuleAnalysisManager &AM) {
+
+  PreservedAnalyses PA;
+
+  DXILResourceBindingInfo &DRBI = AM.getResult<DXILResourceBindingAnalysis>(M);
+  DXILResourceTypeMap &DRTM = AM.getResult<DXILResourceTypeAnalysis>(M);
+  if (DRBI.hasImplicitBinding())
+    if (assignBindings(M, DRBI, DRTM))
+      return PA;
----------------
V-FEXrt wrote:

What does it mean to return an "empty" `PA` here?

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