1 files changed, 130 insertions, 62 deletions

diff --git a/llvm/lib/Target/SPIRV/SPIRVLegalizeImplicitBinding.cpp b/llvm/lib/Target/SPIRV/SPIRVLegalizeImplicitBinding.cpp
index 205895e..fc14a03 100644
--- a/llvm/lib/Target/SPIRV/SPIRVLegalizeImplicitBinding.cpp
+++ b/llvm/lib/Target/SPIRV/SPIRVLegalizeImplicitBinding.cpp
@@ -39,6 +39,10 @@ private:
   void collectBindingInfo(Module &M);
   uint32_t getAndReserveFirstUnusedBinding(uint32_t DescSet);
   void replaceImplicitBindingCalls(Module &M);
+  void replaceResourceHandleCall(Module &M, CallInst *OldCI,
+                                 uint32_t NewBinding);
+  void replaceCounterHandleCall(Module &M, CallInst *OldCI,
+                                uint32_t NewBinding);
   void verifyUniqueOrderIdPerResource(SmallVectorImpl<CallInst *> &Calls);
 
   // A map from descriptor set to a bit vector of used binding numbers.
@@ -56,64 +60,93 @@ struct BindingInfoCollector : public InstVisitor<BindingInfoCollector> {
       : UsedBindings(UsedBindings), ImplicitBindingCalls(ImplicitBindingCalls) {
   }
 
+  void 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);
+  }
+
   void visitCallInst(CallInst &CI) {
     if (CI.getIntrinsicID() == Intrinsic::spv_resource_handlefrombinding) {
       const uint32_t DescSet =
           cast<ConstantInt>(CI.getArgOperand(0))->getZExtValue();
       const uint32_t Binding =
           cast<ConstantInt>(CI.getArgOperand(1))->getZExtValue();
-
-      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);
+      addBinding(DescSet, Binding);
     } else if (CI.getIntrinsicID() ==
                Intrinsic::spv_resource_handlefromimplicitbinding) {
       ImplicitBindingCalls.push_back(&CI);
+    } else if (CI.getIntrinsicID() ==
+               Intrinsic::spv_resource_counterhandlefrombinding) {
+      const uint32_t DescSet =
+          cast<ConstantInt>(CI.getArgOperand(2))->getZExtValue();
+      const uint32_t Binding =
+          cast<ConstantInt>(CI.getArgOperand(1))->getZExtValue();
+      addBinding(DescSet, Binding);
+    } else if (CI.getIntrinsicID() ==
+               Intrinsic::spv_resource_counterhandlefromimplicitbinding) {
+      ImplicitBindingCalls.push_back(&CI);
     }
   }
 };
 
+static uint32_t getOrderId(const CallInst *CI) {
+  uint32_t OrderIdArgIdx = 0;
+  switch (CI->getIntrinsicID()) {
+  case Intrinsic::spv_resource_handlefromimplicitbinding:
+    OrderIdArgIdx = 0;
+    break;
+  case Intrinsic::spv_resource_counterhandlefromimplicitbinding:
+    OrderIdArgIdx = 1;
+    break;
+  default:
+    llvm_unreachable("CallInst is not an implicit binding intrinsic");
+  }
+  return cast<ConstantInt>(CI->getArgOperand(OrderIdArgIdx))->getZExtValue();
+}
+
+static uint32_t getDescSet(const CallInst *CI) {
+  uint32_t DescSetArgIdx;
+  switch (CI->getIntrinsicID()) {
+  case Intrinsic::spv_resource_handlefromimplicitbinding:
+  case Intrinsic::spv_resource_handlefrombinding:
+    DescSetArgIdx = 1;
+    break;
+  case Intrinsic::spv_resource_counterhandlefromimplicitbinding:
+  case Intrinsic::spv_resource_counterhandlefrombinding:
+    DescSetArgIdx = 2;
+    break;
+  default:
+    llvm_unreachable("CallInst is not an implicit binding intrinsic");
+  }
+  return cast<ConstantInt>(CI->getArgOperand(DescSetArgIdx))->getZExtValue();
+}
+
 void SPIRVLegalizeImplicitBinding::collectBindingInfo(Module &M) {
   BindingInfoCollector InfoCollector(UsedBindings, ImplicitBindingCalls);
   InfoCollector.visit(M);
 
   // Sort the collected calls by their order ID.
-  std::sort(
-      ImplicitBindingCalls.begin(), ImplicitBindingCalls.end(),
-      [](const CallInst *A, const CallInst *B) {
-        const uint32_t OrderIdArgIdx = 0;
-        const uint32_t OrderA =
-            cast<ConstantInt>(A->getArgOperand(OrderIdArgIdx))->getZExtValue();
-        const uint32_t OrderB =
-            cast<ConstantInt>(B->getArgOperand(OrderIdArgIdx))->getZExtValue();
-        return OrderA < OrderB;
-      });
+  std::sort(ImplicitBindingCalls.begin(), ImplicitBindingCalls.end(),
+            [](const CallInst *A, const CallInst *B) {
+              return getOrderId(A) < getOrderId(B);
+            });
 }
 
 void SPIRVLegalizeImplicitBinding::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 OrderIdArgIdx = 0;
-    const uint32_t DescSetArgIdx = 1;
-    const uint32_t OrderA =
-        cast<ConstantInt>(Calls[i - 1]->getArgOperand(OrderIdArgIdx))
-            ->getZExtValue();
-    const uint32_t OrderB =
-        cast<ConstantInt>(Calls[i]->getArgOperand(OrderIdArgIdx))
-            ->getZExtValue();
+    const uint32_t OrderA = getOrderId(Calls[i - 1]);
+    const uint32_t OrderB = getOrderId(Calls[i]);
     if (OrderA == OrderB) {
-      const uint32_t DescSetA =
-          cast<ConstantInt>(Calls[i - 1]->getArgOperand(DescSetArgIdx))
-              ->getZExtValue();
-      const uint32_t DescSetB =
-          cast<ConstantInt>(Calls[i]->getArgOperand(DescSetArgIdx))
-              ->getZExtValue();
+      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");
@@ -144,36 +177,26 @@ void SPIRVLegalizeImplicitBinding::replaceImplicitBindingCalls(Module &M) {
   uint32_t lastBindingNumber = -1;
 
   for (CallInst *OldCI : ImplicitBindingCalls) {
-    IRBuilder<> Builder(OldCI);
-    const uint32_t OrderId =
-        cast<ConstantInt>(OldCI->getArgOperand(0))->getZExtValue();
-    const uint32_t DescSet =
-        cast<ConstantInt>(OldCI->getArgOperand(1))->getZExtValue();
-
-    // Reuse an existing binding for this order ID, if one was already assigned.
-    // Otherwise, assign a new binding.
-    const uint32_t NewBinding = (lastOrderId == OrderId)
-                                    ? lastBindingNumber
-                                    : getAndReserveFirstUnusedBinding(DescSet);
-    lastOrderId = OrderId;
-    lastBindingNumber = NewBinding;
-
-    SmallVector<Value *, 8> Args;
-    Args.push_back(Builder.getInt32(DescSet));
-    Args.push_back(Builder.getInt32(NewBinding));
-
-    // Copy the remaining arguments from the old call.
-    for (uint32_t i = 2; i < OldCI->arg_size(); ++i) {
-      Args.push_back(OldCI->getArgOperand(i));
+    const uint32_t OrderId = getOrderId(OldCI);
+    uint32_t BindingNumber;
+    if (OrderId == lastOrderId) {
+      BindingNumber = lastBindingNumber;
+    } else {
+      const uint32_t DescSet = getDescSet(OldCI);
+      BindingNumber = getAndReserveFirstUnusedBinding(DescSet);
     }
 
-    Function *NewFunc = Intrinsic::getOrInsertDeclaration(
-        &M, Intrinsic::spv_resource_handlefrombinding, OldCI->getType());
-    CallInst *NewCI = Builder.CreateCall(NewFunc, Args);
-    NewCI->setCallingConv(OldCI->getCallingConv());
-
-    OldCI->replaceAllUsesWith(NewCI);
-    OldCI->eraseFromParent();
+    if (OldCI->getIntrinsicID() ==
+        Intrinsic::spv_resource_handlefromimplicitbinding) {
+      replaceResourceHandleCall(M, OldCI, BindingNumber);
+    } else {
+      assert(OldCI->getIntrinsicID() ==
+                 Intrinsic::spv_resource_counterhandlefromimplicitbinding &&
+             "Unexpected implicit binding intrinsic");
+      replaceCounterHandleCall(M, OldCI, BindingNumber);
+    }
+    lastOrderId = OrderId;
+    lastBindingNumber = BindingNumber;
   }
 }
 
@@ -196,4 +219,49 @@ INITIALIZE_PASS(SPIRVLegalizeImplicitBinding, "legalize-spirv-implicit-binding",
 
 ModulePass *llvm::createSPIRVLegalizeImplicitBindingPass() {
   return new SPIRVLegalizeImplicitBinding();
-}
\ No newline at end of file
+}
+
+void SPIRVLegalizeImplicitBinding::replaceResourceHandleCall(
+    Module &M, CallInst *OldCI, uint32_t NewBinding) {
+  IRBuilder<> Builder(OldCI);
+  const uint32_t DescSet =
+      cast<ConstantInt>(OldCI->getArgOperand(1))->getZExtValue();
+
+  SmallVector<Value *, 8> Args;
+  Args.push_back(Builder.getInt32(DescSet));
+  Args.push_back(Builder.getInt32(NewBinding));
+
+  // Copy the remaining arguments from the old call.
+  for (uint32_t i = 2; i < OldCI->arg_size(); ++i) {
+    Args.push_back(OldCI->getArgOperand(i));
+  }
+
+  Function *NewFunc = Intrinsic::getOrInsertDeclaration(
+      &M, Intrinsic::spv_resource_handlefrombinding, OldCI->getType());
+  CallInst *NewCI = Builder.CreateCall(NewFunc, Args);
+  NewCI->setCallingConv(OldCI->getCallingConv());
+
+  OldCI->replaceAllUsesWith(NewCI);
+  OldCI->eraseFromParent();
+}
+
+void SPIRVLegalizeImplicitBinding::replaceCounterHandleCall(
+    Module &M, CallInst *OldCI, uint32_t NewBinding) {
+  IRBuilder<> Builder(OldCI);
+  const uint32_t DescSet =
+      cast<ConstantInt>(OldCI->getArgOperand(2))->getZExtValue();
+
+  SmallVector<Value *, 8> Args;
+  Args.push_back(OldCI->getArgOperand(0));
+  Args.push_back(Builder.getInt32(NewBinding));
+  Args.push_back(Builder.getInt32(DescSet));
+
+  Type *Tys[] = {OldCI->getType(), OldCI->getArgOperand(0)->getType()};
+  Function *NewFunc = Intrinsic::getOrInsertDeclaration(
+      &M, Intrinsic::spv_resource_counterhandlefrombinding, Tys);
+  CallInst *NewCI = Builder.CreateCall(NewFunc, Args);
+  NewCI->setCallingConv(OldCI->getCallingConv());
+
+  OldCI->replaceAllUsesWith(NewCI);
+  OldCI->eraseFromParent();
+}