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// and @llvm.spv.resource.handlefromheap intrinsics by replacing them with a
12// call to @llvm.spv.resource.handlefrombinding.
13// It also replaces any @llvm.spv.resource.counterhandlefromimplicitbinding and
14// @llvm.spv.resource.counterhandlefromheap intrinsics with calls to
15// @llvm.spv.resource.counterhandlefrombinding.
16//
17//===----------------------------------------------------------------------===//
18
19#include "SPIRV.h"
20#include "llvm/ADT/BitVector.h"
21#include "llvm/ADT/STLExtras.h"
22#include "llvm/ADT/SmallString.h"
23#include "llvm/ADT/SmallVector.h"
24#include "llvm/IR/IRBuilder.h"
25#include "llvm/IR/InstVisitor.h"
26#include "llvm/IR/Intrinsics.h"
27#include "llvm/IR/IntrinsicsSPIRV.h"
28#include "llvm/IR/Module.h"
29#include "llvm/Pass.h"
30#include "llvm/Support/raw_ostream.h"
31#include <vector>
32
33using namespace llvm;
34
35namespace {
36class SPIRVLegalizeResourceBindingImpl {
37public:
38 bool runOnModule(Module &M);
39
40private:
41 void collectBindingInfo(Module &M);
42 uint32_t getAndReserveFirstUnusedBinding(uint32_t DescSet);
43 bool replaceImplicitBindingCalls(Module &M);
44 bool replaceHeapBindingCalls(Module &M);
45
46 // A map from descriptor set to a bit vector of used binding numbers.
47 std::vector<BitVector> UsedBindings;
48
49 // Set to true by collectBindingInfo() if there are possibly any implicit
50 // binding or heap binding calls in the module (if the module contains a
51 // declaration of implicit binding or heap intrinsic).
52 bool MayHaveImplicitBindings = false;
53 bool MayHaveHeapBindings = false;
54};
55
56class SPIRVLegalizeResourceBindingLegacy : public ModulePass {
57public:
58 static char ID;
59 SPIRVLegalizeResourceBindingLegacy() : ModulePass(ID) {}
60 StringRef getPassName() const override {
61 return "SPIRV Legalize Resource Binding";
62 }
63 bool runOnModule(Module &M) override {
64 return SPIRVLegalizeResourceBindingImpl().runOnModule(M);
65 }
66};
67
68static uint32_t getDescSet(const CallInst *CI) {
69 uint32_t DescSetArgIdx;
70 switch (CI->getIntrinsicID()) {
71 case Intrinsic::spv_resource_handlefromimplicitbinding:
72 DescSetArgIdx = 1;
73 break;
74 case Intrinsic::spv_resource_counterhandlefromimplicitbinding:
75 DescSetArgIdx = 2;
76 break;
77 default:
78 llvm_unreachable("CallInst is not an implicit binding intrinsic");
79 }
80 return cast<ConstantInt>(Val: CI->getArgOperand(i: DescSetArgIdx))->getZExtValue();
81}
82
83// Collect all of the bindings used by llvm.spv.resource.handlefrombinding
84// and llvm.spv.resource.counterhandlefrombinding calls. Also check if there
85// are any implicit binding calls.
86void SPIRVLegalizeResourceBindingImpl::collectBindingInfo(Module &M) {
87
88 auto addBinding = [&](uint32_t DescSet, uint32_t Binding) {
89 if (UsedBindings.size() <= DescSet) {
90 UsedBindings.resize(new_size: DescSet + 1);
91 UsedBindings[DescSet].resize(N: 64);
92 }
93 if (UsedBindings[DescSet].size() <= Binding) {
94 UsedBindings[DescSet].resize(N: 2 * Binding + 1);
95 }
96 UsedBindings[DescSet].set(Binding);
97 };
98
99 auto collectBinding = [&](Function &F, uint32_t ArgDescSetIdx,
100 uint32_t ArgBindingIdx) {
101 for (User *U : F.users()) {
102 if (CallInst *CI = dyn_cast<CallInst>(Val: U)) {
103 const uint32_t DescSet =
104 cast<ConstantInt>(Val: CI->getArgOperand(i: ArgDescSetIdx))->getZExtValue();
105 const uint32_t Binding =
106 cast<ConstantInt>(Val: CI->getArgOperand(i: ArgBindingIdx))->getZExtValue();
107 addBinding(DescSet, Binding);
108 }
109 }
110 };
111
112 for (Function &F : M) {
113 if (!F.isDeclaration())
114 continue;
115
116 switch (F.getIntrinsicID()) {
117 case Intrinsic::spv_resource_handlefrombinding:
118 collectBinding(F, /*ArgDescSetIdx*/ 0, /*ArgBindingIdx*/ 1);
119 break;
120 case Intrinsic::spv_resource_counterhandlefrombinding:
121 collectBinding(F, /*ArgDescSetIdx*/ 1, /*ArgBindingIdx*/ 2);
122 break;
123 case Intrinsic::spv_resource_handlefromimplicitbinding:
124 case Intrinsic::spv_resource_counterhandlefromimplicitbinding:
125 MayHaveImplicitBindings = true;
126 break;
127 case Intrinsic::spv_resource_handlefromheap:
128 case Intrinsic::spv_resource_counterhandlefromheap:
129 MayHaveHeapBindings = true;
130 break;
131 default:
132 break;
133 }
134 }
135}
136
137uint32_t SPIRVLegalizeResourceBindingImpl::getAndReserveFirstUnusedBinding(
138 uint32_t DescSet) {
139 if (UsedBindings.size() <= DescSet) {
140 UsedBindings.resize(new_size: DescSet + 1);
141 UsedBindings[DescSet].resize(N: 64);
142 }
143
144 int NewBinding = UsedBindings[DescSet].find_first_unset();
145 if (NewBinding == -1) {
146 NewBinding = UsedBindings[DescSet].size();
147 UsedBindings[DescSet].resize(N: 2 * NewBinding + 1);
148 }
149
150 UsedBindings[DescSet].set(NewBinding);
151 return NewBinding;
152}
153
154// Replace the implicit binding call with a new call using explicit binding.
155static void replaceWithHandleFromBinding(Module &M, CallInst *CI,
156 uint32_t DescSet, uint32_t Binding,
157 Value *RangeOp, Value *IndexOp,
158 Value *Name) {
159 assert((CI->getIntrinsicID() ==
160 Intrinsic::spv_resource_handlefromimplicitbinding ||
161 CI->getIntrinsicID() == Intrinsic::spv_resource_handlefromheap) &&
162 "unexpected binding intrinsic");
163 IRBuilder<> Builder(CI);
164 Value *DescSetOp = Builder.getInt32(C: DescSet);
165 Value *BindingOp = Builder.getInt32(C: Binding);
166 Function *NewFunc = Intrinsic::getOrInsertDeclaration(
167 M: &M, id: Intrinsic::spv_resource_handlefrombinding, OverloadTys: {CI->getType()});
168 CallInst *NewCI = Builder.CreateCall(
169 Callee: NewFunc, Args: {DescSetOp, BindingOp, RangeOp, IndexOp, Name});
170 NewCI->setCallingConv(CI->getCallingConv());
171 CI->replaceAllUsesWith(V: NewCI);
172 CI->eraseFromParent();
173}
174
175// Replace the implicit counter binding call with a new call using explicit
176// binding.
177static void replaceWithCounterHandleFromBinding(Module &M, CallInst *CI,
178 uint32_t DescSet,
179 uint32_t Binding) {
180 assert(
181 (CI->getIntrinsicID() ==
182 Intrinsic::spv_resource_counterhandlefromimplicitbinding ||
183 CI->getIntrinsicID() == Intrinsic::spv_resource_counterhandlefromheap) &&
184 "unexpected binding intrinsic");
185 IRBuilder<> Builder(CI);
186 Value *DescSetOp = Builder.getInt32(C: DescSet);
187 Value *BindingOp = Builder.getInt32(C: Binding);
188 Value *MainHandle = CI->getArgOperand(i: 0);
189 Type *OverloadTys[] = {CI->getType(), MainHandle->getType()};
190 Function *NewFunc = Intrinsic::getOrInsertDeclaration(
191 M: &M, id: Intrinsic::spv_resource_counterhandlefrombinding, OverloadTys);
192 CallInst *NewCI =
193 Builder.CreateCall(Callee: NewFunc, Args: {MainHandle, DescSetOp, BindingOp});
194 NewCI->setCallingConv(CI->getCallingConv());
195 CI->replaceAllUsesWith(V: NewCI);
196 CI->eraseFromParent();
197}
198
199bool SPIRVLegalizeResourceBindingImpl::replaceImplicitBindingCalls(Module &M) {
200 // Collect all implicit binding calls.
201 SmallVector<std::pair<uint32_t, CallInst *>> IBCalls;
202 bool Changed = false;
203 for (Function &F : M) {
204 if (!F.isDeclaration())
205 continue;
206
207 uint32_t OrderIdIdx;
208 if (F.getIntrinsicID() == Intrinsic::spv_resource_handlefromimplicitbinding)
209 OrderIdIdx = 0;
210 else if (F.getIntrinsicID() ==
211 Intrinsic::spv_resource_counterhandlefromimplicitbinding)
212 OrderIdIdx = 1;
213 else
214 continue;
215
216 for (User *U : F.users()) {
217 if (CallInst *CI = dyn_cast<CallInst>(Val: U)) {
218 ConstantInt *OrderId = cast<ConstantInt>(Val: CI->getArgOperand(i: OrderIdIdx));
219 IBCalls.emplace_back(Args: OrderId->getZExtValue(), Args&: CI);
220 }
221 }
222 }
223
224 if (IBCalls.empty())
225 return false;
226
227 // Sort the collected calls by their order ID.
228 llvm::sort(C&: IBCalls, Comp: llvm::less_first());
229
230 // Assign bindings based on the order ID. Same order ID gets the same binding.
231 // Also make sure that calls with the same order ID have the same descriptor
232 // set.
233 uint32_t LastOrderId = -1;
234 uint32_t LastBinding = -1;
235 uint32_t LastDescSet = -1;
236 for (auto &[OrderId, CI] : IBCalls) {
237 uint32_t Binding;
238 uint32_t DescSet = getDescSet(CI);
239 if (OrderId == LastOrderId) {
240 if (DescSet != LastDescSet)
241 report_fatal_error(reason: "Implicit binding calls with the same order ID must "
242 "have the same descriptor set");
243 Binding = LastBinding;
244 } else {
245 Binding = getAndReserveFirstUnusedBinding(DescSet);
246 }
247
248 // Replace the implicit binding call with an explicit binding call.
249 if (CI->getIntrinsicID() ==
250 Intrinsic::spv_resource_handlefromimplicitbinding)
251 replaceWithHandleFromBinding(M, CI, DescSet, Binding,
252 RangeOp: CI->getArgOperand(i: 2), IndexOp: CI->getArgOperand(i: 3),
253 Name: CI->getArgOperand(i: 4));
254 else
255 replaceWithCounterHandleFromBinding(M, CI, DescSet, Binding);
256 Changed = true;
257
258 LastOrderId = OrderId;
259 LastBinding = Binding;
260 LastDescSet = DescSet;
261 }
262 return Changed;
263}
264
265bool moduleContainsConstantString(Module &M, StringRef Str) {
266 for (GlobalVariable &GV : M.globals()) {
267 if (!GV.hasInitializer())
268 continue;
269 if (ConstantDataArray *CDA =
270 dyn_cast<ConstantDataArray>(Val: GV.getInitializer())) {
271 if (CDA->isString() && CDA->getAsCString() == Str)
272 return true;
273 }
274 }
275 return false;
276}
277
278// Creates unique global variable for the heap name string, making
279// sure that each heap name string is unique within the module.
280GlobalVariable *createHeapNameString(Module &M, StringRef Name) {
281 SmallString<32> GlobalStringName(Name);
282 uint32_t HeapNameLen = Name.size();
283 StringRef HeapName;
284 for (unsigned Suffix = 1;; ++Suffix) {
285 GlobalStringName.append(RHS: ".str");
286 if (!M.getNamedValue(Name: GlobalStringName)) {
287 // Make sure the module does not already have a constant
288 // string with this value.
289 HeapName = GlobalStringName.substr(Start: 0, N: HeapNameLen);
290 if (!moduleContainsConstantString(M, Str: HeapName))
291 break;
292 }
293 GlobalStringName.resize(N: Name.size());
294 raw_svector_ostream(GlobalStringName) << '.' << Suffix;
295 HeapNameLen = GlobalStringName.size();
296 }
297 Constant *Init = ConstantDataArray::getString(Context&: M.getContext(), Initializer: HeapName);
298 GlobalVariable *HeapNameGV = new GlobalVariable(
299 M, Init->getType(), /*isConstant=*/true, GlobalValue::PrivateLinkage,
300 Init, GlobalStringName, /*InsertBefore=*/nullptr,
301 GlobalVariable::NotThreadLocal, /*AddressSpace=*/0);
302 HeapNameGV->setUnnamedAddr(GlobalValue::UnnamedAddr::Global);
303 HeapNameGV->setAlignment(Align(1));
304 return HeapNameGV;
305}
306
307// The SPIR-V backend represents dynamic resources as unbounded resource arrays.
308// This function scans the module for calls to
309// `llvm.spv.resource.handlefromheap` and groups them according to whether they
310// create CBV/SRV/UAV resources or samplers. It also collects calls to
311// `llvm.spv.resource.counterhandlefromheap` intrinsics that form a third group.
312//
313// The function assigns the first available binding to each non-empty heap group
314// in this order: CBV/SRV/UAV resources, samplers, and counters. For each group,
315// it will replace the heap intrinsic calls with their explicit
316// `handlefrombinding` equivalents using the assigned binding.
317//
318// The function does not actually create the unbounded resource-array globals
319// itself. It only assigns a unique name that is shared by all resources
320// belonging to the same heap type; the existing `SPIRVInstructionSelector` will
321// create the globals.
322//
323// Because the CBV/SRV/UAV group can contain resources of different types, these
324// resources must be represented by separate arrays, one for each unique
325// resource type. All of these arrays will use the same binding and therefore
326// overlap.
327bool SPIRVLegalizeResourceBindingImpl::replaceHeapBindingCalls(Module &M) {
328 // First we collect all used heap binding declarations and group them based on
329 // their kind.
330 SmallVector<Function *, 8> CbvSrvUavs;
331 SmallVector<Function *, 8> Samplers;
332 SmallVector<Function *, 8> Counters;
333 bool Changed = false;
334
335 for (Function &F : M) {
336 if (!F.isDeclaration() || F.user_empty())
337 continue;
338
339 if (F.getIntrinsicID() == Intrinsic::spv_resource_handlefromheap) {
340 TargetExtType *ResType = cast<TargetExtType>(Val: F.getReturnType());
341 if (ResType->getName() == "spirv.Sampler")
342 Samplers.emplace_back(Args: &F);
343 else
344 CbvSrvUavs.emplace_back(Args: &F);
345 } else if (F.getIntrinsicID() ==
346 Intrinsic::spv_resource_counterhandlefromheap) {
347 Counters.emplace_back(Args: &F);
348 } else
349 continue;
350 }
351
352 if (CbvSrvUavs.empty() && Samplers.empty() && Counters.empty())
353 return false;
354
355 // Heap resources are always mapped to descriptor set 0 as an unbounded
356 // runtime array.
357 constexpr uint32_t DescSet = 0;
358 Value *Zero =
359 llvm::ConstantInt::get(Ty: llvm::Type::getInt32Ty(C&: M.getContext()), V: 0);
360
361 if (!CbvSrvUavs.empty()) {
362 // For CBV/UAV/SRV resources we need to create a different
363 // ResourceDescriptorHeap name for each unique resource type. They will all
364 // share the same binding and will overlap.
365 uint32_t Binding = getAndReserveFirstUnusedBinding(DescSet);
366 SmallDenseMap<TargetExtType *, GlobalVariable *> ResourceDescriptorHeaps;
367 for (Function *F : CbvSrvUavs) {
368 TargetExtType *ResType = cast<TargetExtType>(Val: F->getReturnType());
369 GlobalVariable *HeapNameGV = nullptr;
370 auto It = ResourceDescriptorHeaps.find(Val: ResType);
371 if (It == ResourceDescriptorHeaps.end()) {
372 HeapNameGV = createHeapNameString(M, Name: "ResourceDescriptorHeap");
373 [[maybe_unused]] auto [InsertedIt, Inserted] =
374 ResourceDescriptorHeaps.try_emplace(Key: ResType, Args&: HeapNameGV);
375 assert(Inserted && "resource heap name already exists");
376 } else {
377 HeapNameGV = It->second;
378 }
379
380 for (User *U : make_early_inc_range(Range: F->users())) {
381 if (CallInst *CI = dyn_cast<CallInst>(Val: U)) {
382 Value *HeapIdx = CI->getArgOperand(i: 0);
383 replaceWithHandleFromBinding(M, CI, DescSet, Binding, RangeOp: Zero, IndexOp: HeapIdx,
384 Name: HeapNameGV);
385 Changed = true;
386 }
387 }
388 F->eraseFromParent();
389 }
390 }
391
392 if (!Samplers.empty()) {
393 uint32_t Binding = getAndReserveFirstUnusedBinding(DescSet);
394 // The Sampler handle type should be the same for all samplers
395 // (target("spirv.Sampler")).
396 [[maybe_unused]] TargetExtType *SamplerHandleType =
397 cast<TargetExtType>(Val: Samplers.front()->getReturnType());
398
399 GlobalVariable *HeapNameGV =
400 createHeapNameString(M, Name: "SamplerDescriptorHeap");
401 for (Function *F : Samplers) {
402 assert(F->getReturnType() == SamplerHandleType &&
403 "sampler handle type mismatch");
404 for (User *U : make_early_inc_range(Range: F->users())) {
405 if (CallInst *CI = dyn_cast<CallInst>(Val: U)) {
406 Value *HeapIdx = CI->getArgOperand(i: 0);
407 replaceWithHandleFromBinding(M, CI, DescSet, Binding, RangeOp: Zero, IndexOp: HeapIdx,
408 Name: HeapNameGV);
409 Changed = true;
410 }
411 }
412 F->eraseFromParent();
413 }
414 }
415
416 if (!Counters.empty()) {
417 uint32_t Binding = getAndReserveFirstUnusedBinding(DescSet);
418 [[maybe_unused]] Type *CounterHandleTy = Counters.front()->getReturnType();
419 for (Function *F : Counters) {
420 // The counter handle type should be the same for all resource types
421 // that have a counter (target("spirv.VulkanBuffer", i32, 12, 1)).
422 assert(F->getReturnType() == CounterHandleTy &&
423 "counter handle type mismatch");
424 for (User *U : make_early_inc_range(Range: F->users())) {
425 if (CallInst *CI = dyn_cast<CallInst>(Val: U)) {
426 replaceWithCounterHandleFromBinding(M, CI, DescSet, Binding);
427 Changed = true;
428 }
429 }
430 F->eraseFromParent();
431 }
432 }
433
434 return Changed;
435}
436
437bool SPIRVLegalizeResourceBindingImpl::runOnModule(Module &M) {
438 collectBindingInfo(M);
439
440 bool Changed = false;
441 if (MayHaveImplicitBindings)
442 Changed |= replaceImplicitBindingCalls(M);
443 if (MayHaveHeapBindings)
444 Changed |= replaceHeapBindingCalls(M);
445
446 return Changed;
447}
448} // namespace
449
450PreservedAnalyses
451SPIRVLegalizeResourceBindingPass::run(Module &M, ModuleAnalysisManager &AM) {
452 return SPIRVLegalizeResourceBindingImpl().runOnModule(M)
453 ? PreservedAnalyses::none()
454 : PreservedAnalyses::all();
455}
456
457char SPIRVLegalizeResourceBindingLegacy::ID = 0;
458
459INITIALIZE_PASS(SPIRVLegalizeResourceBindingLegacy,
460 "legalize-spirv-resource-binding",
461 "Legalize SPIR-V resource bindings", false, false)
462
463ModulePass *llvm::createSPIRVLegalizeResourceBindingPass() {
464 return new SPIRVLegalizeResourceBindingLegacy();
465}
466