[llvm-branch-commits] [llvm] [SPIRV] Refactor implicit binding legalization (PR #221662)

[Helena Kotas via llvm-branch-commits]

================
@@ -182,43 +139,121 @@ uint32_t SPIRVLegalizeImplicitBindingImpl::getAndReserveFirstUnusedBinding(
   return NewBinding;
 }
 
-void SPIRVLegalizeImplicitBindingImpl::replaceImplicitBindingCalls(Module &M) {
-  uint32_t lastOrderId = -1;
-  uint32_t lastBindingNumber = -1;
+// Replace the implicit binding call with a new call using explicit binding.
+static void replaceWithHandleFromBinding(Module &M, CallInst *CI,
+                                         uint32_t DescSet, uint32_t Binding,
+                                         Value *IndexOp, Value *RangeOp,
+                                         Value *Name) {
+  IRBuilder<> Builder(CI);
+  Value *DescSetOp = Builder.getInt32(DescSet);
+  Value *BindingOp = Builder.getInt32(Binding);
+  Function *NewFunc = Intrinsic::getOrInsertDeclaration(
+      &M, Intrinsic::spv_resource_handlefrombinding, {CI->getType()});
+  CallInst *NewCI = Builder.CreateCall(
+      NewFunc, {DescSetOp, BindingOp, IndexOp, RangeOp, Name});
+  NewCI->setCallingConv(CI->getCallingConv());
+  CI->replaceAllUsesWith(NewCI);
+  CI->eraseFromParent();
+}
 
-  for (CallInst *OldCI : ImplicitBindingCalls) {
-    const uint32_t OrderId = getOrderId(OldCI);
-    uint32_t BindingNumber;
-    if (OrderId == lastOrderId) {
-      BindingNumber = lastBindingNumber;
-    } else {
-      const uint32_t DescSet = getDescSet(OldCI);
-      BindingNumber = getAndReserveFirstUnusedBinding(DescSet);
+// Replace the implicit counter binding call with a new call using explicit
+// binding.
+static void replaceWithCounterHandleFromBinding(Module &M, CallInst *CI,
+                                                Value *MainHandle,
+                                                uint32_t DescSet,
+                                                uint32_t Binding) {
+
+  assert(CI->getIntrinsicID() ==
+             Intrinsic::spv_resource_counterhandlefromimplicitbinding &&
+         "unexpected implicit binding intrinsic");
+  IRBuilder<> Builder(CI);
+  Value *DescSetOp = Builder.getInt32(DescSet);
+  Value *BindingOp = Builder.getInt32(Binding);
+  Type *OverloadTys[] = {CI->getType(), CI->getArgOperand(0)->getType()};
+  Function *NewFunc = Intrinsic::getOrInsertDeclaration(
+      &M, Intrinsic::spv_resource_counterhandlefrombinding, OverloadTys);
+  CallInst *NewCI =
+      Builder.CreateCall(NewFunc, {MainHandle, DescSetOp, BindingOp});
+  NewCI->setCallingConv(CI->getCallingConv());
+  CI->replaceAllUsesWith(NewCI);
+  CI->eraseFromParent();
+}
+
+bool SPIRVLegalizeImplicitBindingImpl::replaceImplicitBindingCalls(Module &M) {
+  // Collect all implicit binding calls.
+  SmallVector<std::pair<uint32_t, CallInst *>> IBCalls;
+  bool Changed = false;
+  for (Function &F : M) {
+    if (!F.isDeclaration())
+      continue;
+
+    uint32_t OrderIdIdx;
+    if (F.getIntrinsicID() == Intrinsic::spv_resource_handlefromimplicitbinding)
+      OrderIdIdx = 0;
+    else if (F.getIntrinsicID() ==
+             Intrinsic::spv_resource_counterhandlefromimplicitbinding)
+      OrderIdIdx = 1;
+    else
+      continue;
+
+    for (User *U : F.users()) {
+      if (CallInst *CI = dyn_cast<CallInst>(U)) {
+        ConstantInt *OrderId = cast<ConstantInt>(CI->getArgOperand(OrderIdIdx));
+        IBCalls.emplace_back(OrderId->getZExtValue(), CI);
+      }
     }
+  }
 
-    if (OldCI->getIntrinsicID() ==
-        Intrinsic::spv_resource_handlefromimplicitbinding) {
-      replaceResourceHandleCall(M, OldCI, BindingNumber);
+  // Sort the collected calls by their order ID.
+  llvm::sort(IBCalls, [](const std::pair<uint32_t, CallInst *> &A,
+                         const std::pair<uint32_t, CallInst *> &B) {
+    return A.first < B.first;
+  });
+
+  // Assign bindings based on the order ID. Same order ID gets the same binding.
+  // Also make sure that calls with the same order ID have the same descriptor
+  // set.
+  uint32_t LastOrderId = -1;
+  uint32_t LastBinding = -1;
+  uint32_t LastDescSet = -1;
+  for (auto &[OrderId, CI] : IBCalls) {
+    uint32_t Binding;
+    uint32_t DescSet = getDescSet(CI);
+    if (OrderId == LastOrderId) {
+      if (DescSet != LastDescSet)
+        report_fatal_error("Implicit binding calls with the same order ID must "
+                           "have the same descriptor set");
+      Binding = LastBinding;
     } else {
-      assert(OldCI->getIntrinsicID() ==
-                 Intrinsic::spv_resource_counterhandlefromimplicitbinding &&
-             "Unexpected implicit binding intrinsic");
-      replaceCounterHandleCall(M, OldCI, BindingNumber);
+      Binding = getAndReserveFirstUnusedBinding(DescSet);
     }
-    lastOrderId = OrderId;
-    lastBindingNumber = BindingNumber;
+
+    // Replace the implicit binding call with an explicit binding call.
+    if (CI->getIntrinsicID() ==
+        Intrinsic::spv_resource_handlefromimplicitbinding)
+      replaceWithHandleFromBinding(M, CI, DescSet, Binding,
+                                   CI->getArgOperand(2), CI->getArgOperand(3),
+                                   CI->getArgOperand(4));
+    else
+      replaceWithCounterHandleFromBinding(M, CI, CI->getArgOperand(0), DescSet,
+                                          Binding);
+    Changed = true;
+
+    LastOrderId = OrderId;
+    LastBinding = Binding;
+    LastDescSet = DescSet;
   }
+  return Changed;
 }
 
 bool SPIRVLegalizeImplicitBindingImpl::runOnModule(Module &M) {
   collectBindingInfo(M);
-  if (ImplicitBindingCalls.empty()) {
-    return false;
-  }
-  verifyUniqueOrderIdPerResource(ImplicitBindingCalls);
 
-  replaceImplicitBindingCalls(M);
-  return true;
+  bool Changed = false;
+  if (HasImplicitBinding)
+    Changed |= replaceImplicitBindingCalls(M);
+
+  return Changed;
----------------
hekota wrote:

Since we are already iterating though all the declarations in `collectBindingInfo`, I believe it is worth to also gather some hints that would allow for early exit, like if the module has any implicit bindings. That way we can skip `replaceImplicitBindingCalls` that would iterate over the declarations again. I will update the name to `MayHaveImplicitBindings` to reflect the case when the declaration has no uses.

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