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

[Helena Kotas 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);
----------------
hekota wrote:

`replaceAllUsesWith` is updating all places that use the value returned from the original call with the value returned from the new call. So if the handle returned from the original call is used in multiple calls like `llvm.dx.resource.getpointer` or `llvm.dx.resource.updatecounter`, it will replace the handle argument in all of these calls with the new value returned from the new call.

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