1//===- SPIRVLegalizeResourceBinding.cpp - Legalize resource bindings ----*- C++
2//-*-===//
3//
4// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
5// See https://llvm.org/LICENSE.txt for license information.
6// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
7//
8//===----------------------------------------------------------------------===//
9//
10// This pass legalizes the @llvm.spv.resource.handlefromimplicitbinding
11// intrinsic by replacing it with a call to
12// @llvm.spv.resource.handlefrombinding.
13//
14//===----------------------------------------------------------------------===//
15
16#include "SPIRV.h"
17#include "llvm/ADT/BitVector.h"
18#include "llvm/ADT/STLExtras.h"
19#include "llvm/ADT/SmallVector.h"
20#include "llvm/IR/IRBuilder.h"
21#include "llvm/IR/InstVisitor.h"
22#include "llvm/IR/Intrinsics.h"
23#include "llvm/IR/IntrinsicsSPIRV.h"
24#include "llvm/IR/Module.h"
25#include "llvm/Pass.h"
26#include <vector>
27
28using namespace llvm;
29
30namespace {
31class SPIRVLegalizeResourceBindingImpl {
32public:
33 bool runOnModule(Module &M);
34
35private:
36 void collectBindingInfo(Module &M);
37 uint32_t getAndReserveFirstUnusedBinding(uint32_t DescSet);
38 bool replaceImplicitBindingCalls(Module &M);
39
40 // A map from descriptor set to a bit vector of used binding numbers.
41 std::vector<BitVector> UsedBindings;
42
43 // Set to true by collectBindingInfo() if there are any implicit binding
44 // declarations in the module.
45 bool MayHaveImplicitBindings = false;
46};
47
48class SPIRVLegalizeResourceBindingLegacy : public ModulePass {
49public:
50 static char ID;
51 SPIRVLegalizeResourceBindingLegacy() : ModulePass(ID) {}
52 StringRef getPassName() const override {
53 return "SPIRV Legalize Resource Binding";
54 }
55 bool runOnModule(Module &M) override {
56 return SPIRVLegalizeResourceBindingImpl().runOnModule(M);
57 }
58};
59
60static uint32_t getDescSet(const CallInst *CI) {
61 uint32_t DescSetArgIdx;
62 switch (CI->getIntrinsicID()) {
63 case Intrinsic::spv_resource_handlefromimplicitbinding:
64 DescSetArgIdx = 1;
65 break;
66 case Intrinsic::spv_resource_counterhandlefromimplicitbinding:
67 DescSetArgIdx = 2;
68 break;
69 default:
70 llvm_unreachable("CallInst is not an implicit binding intrinsic");
71 }
72 return cast<ConstantInt>(Val: CI->getArgOperand(i: DescSetArgIdx))->getZExtValue();
73}
74
75// Collect all of the bindings used by llvm.spv.resource.handlefrombinding
76// and llvm.spv.resource.counterhandlefrombinding calls. Also check if there
77// are any implicit binding calls.
78void SPIRVLegalizeResourceBindingImpl::collectBindingInfo(Module &M) {
79
80 auto addBinding = [&](uint32_t DescSet, uint32_t Binding) {
81 if (UsedBindings.size() <= DescSet) {
82 UsedBindings.resize(new_size: DescSet + 1);
83 UsedBindings[DescSet].resize(N: 64);
84 }
85 if (UsedBindings[DescSet].size() <= Binding) {
86 UsedBindings[DescSet].resize(N: 2 * Binding + 1);
87 }
88 UsedBindings[DescSet].set(Binding);
89 };
90
91 auto collectBinding = [&](Function &F, uint32_t ArgDescSetIdx,
92 uint32_t ArgBindingIdx) {
93 for (User *U : F.users()) {
94 if (CallInst *CI = dyn_cast<CallInst>(Val: U)) {
95 const uint32_t DescSet =
96 cast<ConstantInt>(Val: CI->getArgOperand(i: ArgDescSetIdx))->getZExtValue();
97 const uint32_t Binding =
98 cast<ConstantInt>(Val: CI->getArgOperand(i: ArgBindingIdx))->getZExtValue();
99 addBinding(DescSet, Binding);
100 }
101 }
102 };
103
104 for (Function &F : M) {
105 if (!F.isDeclaration())
106 continue;
107
108 switch (F.getIntrinsicID()) {
109 case Intrinsic::spv_resource_handlefrombinding:
110 collectBinding(F, /*ArgDescSetIdx*/ 0, /*ArgBindingIdx*/ 1);
111 break;
112 case Intrinsic::spv_resource_counterhandlefrombinding:
113 collectBinding(F, /*ArgDescSetIdx*/ 1, /*ArgBindingIdx*/ 2);
114 break;
115 case Intrinsic::spv_resource_handlefromimplicitbinding:
116 case Intrinsic::spv_resource_counterhandlefromimplicitbinding:
117 MayHaveImplicitBindings = true;
118 break;
119 default:
120 break;
121 }
122 }
123}
124
125uint32_t SPIRVLegalizeResourceBindingImpl::getAndReserveFirstUnusedBinding(
126 uint32_t DescSet) {
127 if (UsedBindings.size() <= DescSet) {
128 UsedBindings.resize(new_size: DescSet + 1);
129 UsedBindings[DescSet].resize(N: 64);
130 }
131
132 int NewBinding = UsedBindings[DescSet].find_first_unset();
133 if (NewBinding == -1) {
134 NewBinding = UsedBindings[DescSet].size();
135 UsedBindings[DescSet].resize(N: 2 * NewBinding + 1);
136 }
137
138 UsedBindings[DescSet].set(NewBinding);
139 return NewBinding;
140}
141
142// Replace the implicit binding call with a new call using explicit binding.
143static void replaceWithHandleFromBinding(Module &M, CallInst *CI,
144 uint32_t DescSet, uint32_t Binding,
145 Value *IndexOp, Value *RangeOp,
146 Value *Name) {
147 assert(CI->getIntrinsicID() ==
148 Intrinsic::spv_resource_handlefromimplicitbinding &&
149 "unexpected implicit binding intrinsic");
150 IRBuilder<> Builder(CI);
151 Value *DescSetOp = Builder.getInt32(C: DescSet);
152 Value *BindingOp = Builder.getInt32(C: Binding);
153 Function *NewFunc = Intrinsic::getOrInsertDeclaration(
154 M: &M, id: Intrinsic::spv_resource_handlefrombinding, OverloadTys: {CI->getType()});
155 CallInst *NewCI = Builder.CreateCall(
156 Callee: NewFunc, Args: {DescSetOp, BindingOp, IndexOp, RangeOp, Name});
157 NewCI->setCallingConv(CI->getCallingConv());
158 CI->replaceAllUsesWith(V: NewCI);
159 CI->eraseFromParent();
160}
161
162// Replace the implicit counter binding call with a new call using explicit
163// binding.
164static void replaceWithCounterHandleFromBinding(Module &M, CallInst *CI,
165 uint32_t DescSet,
166 uint32_t Binding) {
167 assert(CI->getIntrinsicID() ==
168 Intrinsic::spv_resource_counterhandlefromimplicitbinding &&
169 "unexpected implicit binding intrinsic");
170 IRBuilder<> Builder(CI);
171 Value *DescSetOp = Builder.getInt32(C: DescSet);
172 Value *BindingOp = Builder.getInt32(C: Binding);
173 Value *MainHandle = CI->getArgOperand(i: 0);
174 Type *OverloadTys[] = {CI->getType(), MainHandle->getType()};
175 Function *NewFunc = Intrinsic::getOrInsertDeclaration(
176 M: &M, id: Intrinsic::spv_resource_counterhandlefrombinding, OverloadTys);
177 CallInst *NewCI =
178 Builder.CreateCall(Callee: NewFunc, Args: {MainHandle, DescSetOp, BindingOp});
179 NewCI->setCallingConv(CI->getCallingConv());
180 CI->replaceAllUsesWith(V: NewCI);
181 CI->eraseFromParent();
182}
183
184bool SPIRVLegalizeResourceBindingImpl::replaceImplicitBindingCalls(Module &M) {
185 // Collect all implicit binding calls.
186 SmallVector<std::pair<uint32_t, CallInst *>> IBCalls;
187 bool Changed = false;
188 for (Function &F : M) {
189 if (!F.isDeclaration())
190 continue;
191
192 uint32_t OrderIdIdx;
193 if (F.getIntrinsicID() == Intrinsic::spv_resource_handlefromimplicitbinding)
194 OrderIdIdx = 0;
195 else if (F.getIntrinsicID() ==
196 Intrinsic::spv_resource_counterhandlefromimplicitbinding)
197 OrderIdIdx = 1;
198 else
199 continue;
200
201 for (User *U : F.users()) {
202 if (CallInst *CI = dyn_cast<CallInst>(Val: U)) {
203 ConstantInt *OrderId = cast<ConstantInt>(Val: CI->getArgOperand(i: OrderIdIdx));
204 IBCalls.emplace_back(Args: OrderId->getZExtValue(), Args&: CI);
205 }
206 }
207 }
208
209 if (IBCalls.empty())
210 return false;
211
212 // Sort the collected calls by their order ID.
213 llvm::sort(C&: IBCalls, Comp: llvm::less_first());
214
215 // Assign bindings based on the order ID. Same order ID gets the same binding.
216 // Also make sure that calls with the same order ID have the same descriptor
217 // set.
218 uint32_t LastOrderId = -1;
219 uint32_t LastBinding = -1;
220 uint32_t LastDescSet = -1;
221 for (auto &[OrderId, CI] : IBCalls) {
222 uint32_t Binding;
223 uint32_t DescSet = getDescSet(CI);
224 if (OrderId == LastOrderId) {
225 if (DescSet != LastDescSet)
226 report_fatal_error(reason: "Implicit binding calls with the same order ID must "
227 "have the same descriptor set");
228 Binding = LastBinding;
229 } else {
230 Binding = getAndReserveFirstUnusedBinding(DescSet);
231 }
232
233 // Replace the implicit binding call with an explicit binding call.
234 if (CI->getIntrinsicID() ==
235 Intrinsic::spv_resource_handlefromimplicitbinding)
236 replaceWithHandleFromBinding(M, CI, DescSet, Binding,
237 IndexOp: CI->getArgOperand(i: 2), RangeOp: CI->getArgOperand(i: 3),
238 Name: CI->getArgOperand(i: 4));
239 else
240 replaceWithCounterHandleFromBinding(M, CI, DescSet, Binding);
241 Changed = true;
242
243 LastOrderId = OrderId;
244 LastBinding = Binding;
245 LastDescSet = DescSet;
246 }
247 return Changed;
248}
249
250bool SPIRVLegalizeResourceBindingImpl::runOnModule(Module &M) {
251 collectBindingInfo(M);
252
253 bool Changed = false;
254 if (MayHaveImplicitBindings)
255 Changed |= replaceImplicitBindingCalls(M);
256
257 return Changed;
258}
259} // namespace
260
261PreservedAnalyses
262SPIRVLegalizeResourceBindingPass::run(Module &M, ModuleAnalysisManager &AM) {
263 return SPIRVLegalizeResourceBindingImpl().runOnModule(M)
264 ? PreservedAnalyses::none()
265 : PreservedAnalyses::all();
266}
267
268char SPIRVLegalizeResourceBindingLegacy::ID = 0;
269
270INITIALIZE_PASS(SPIRVLegalizeResourceBindingLegacy,
271 "legalize-spirv-resource-binding",
272 "Legalize SPIR-V resource bindings", false, false)
273
274ModulePass *llvm::createSPIRVLegalizeResourceBindingPass() {
275 return new SPIRVLegalizeResourceBindingLegacy();
276}
277