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

[Deric C. 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;
+  });
----------------
Icohedron wrote:

[SUGGESTION] This can be shortened with an llvm helper:
```suggestion
  llvm::sort(IBCalls, llvm::less_first());
```

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