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

[Deric C. via llvm-branch-commits]

================
@@ -138,30 +72,53 @@ static uint32_t getDescSet(const CallInst *CI) {
   return cast<ConstantInt>(CI->getArgOperand(DescSetArgIdx))->getZExtValue();
 }
 
+// Collect all of the bindings used by llvm.spv.resource.handlefrombinding
+// and llvm.spv.resource.counterhandlefrombinding calls. Also check if there
+// are any implicit binding calls.
 void SPIRVLegalizeImplicitBindingImpl::collectBindingInfo(Module &M) {
-  BindingInfoCollector InfoCollector(UsedBindings, ImplicitBindingCalls);
-  InfoCollector.visit(M);
-
-  // Sort the collected calls by their order ID.
-  llvm::sort(ImplicitBindingCalls, [](const CallInst *A, const CallInst *B) {
-    return getOrderId(A) < getOrderId(B);
-  });
-}
 
-void SPIRVLegalizeImplicitBindingImpl::verifyUniqueOrderIdPerResource(
-    SmallVectorImpl<CallInst *> &Calls) {
-  // Check that the order Id is unique per resource.
-  for (uint32_t i = 1; i < Calls.size(); ++i) {
-    const uint32_t OrderA = getOrderId(Calls[i - 1]);
-    const uint32_t OrderB = getOrderId(Calls[i]);
-    if (OrderA == OrderB) {
-      const uint32_t DescSetA = getDescSet(Calls[i - 1]);
-      const uint32_t DescSetB = getDescSet(Calls[i]);
-      if (DescSetA != DescSetB) {
-        report_fatal_error("Implicit binding calls with the same order ID must "
-                           "have the same descriptor set");
+  auto addBinding = [&](uint32_t DescSet, uint32_t Binding) {
+    if (UsedBindings.size() <= DescSet) {
+      UsedBindings.resize(DescSet + 1);
+      UsedBindings[DescSet].resize(64);
+    }
+    if (UsedBindings[DescSet].size() <= Binding) {
+      UsedBindings[DescSet].resize(2 * Binding + 1);
+    }
+    UsedBindings[DescSet].set(Binding);
+  };
+
+  auto collectBinding = [&](Function &F, uint32_t ArgDescSetIdx,
+                            uint32_t ArgBindingIdx) {
+    for (User *U : F.users()) {
+      if (CallInst *CI = dyn_cast<CallInst>(U)) {
+        const uint32_t DescSet =
+            cast<ConstantInt>(CI->getArgOperand(ArgDescSetIdx))->getZExtValue();
+        const uint32_t Binding =
+            cast<ConstantInt>(CI->getArgOperand(ArgBindingIdx))->getZExtValue();
+        addBinding(DescSet, Binding);
       }
     }
+  };
+
+  for (Function &F : M.functions()) {
----------------
Icohedron wrote:

[NIT] In `replaceImplicitBindingCalls` a similar loop iterates over `M` instead of `M.functions()`. They are equivalent. Prefer keeping them consistent.
```suggestion
  for (Function &F : M) {
```

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