1 | //===-- SPIRVRegularizer.cpp - regularize IR for SPIR-V ---------*- C++ -*-===// |
2 | // |
3 | // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
4 | // See https://llvm.org/LICENSE.txt for license information. |
5 | // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
6 | // |
7 | //===----------------------------------------------------------------------===// |
8 | // |
9 | // This pass implements regularization of LLVM IR for SPIR-V. The prototype of |
10 | // the pass was taken from SPIRV-LLVM translator. |
11 | // |
12 | //===----------------------------------------------------------------------===// |
13 | |
14 | #include "SPIRV.h" |
15 | #include "llvm/Demangle/Demangle.h" |
16 | #include "llvm/IR/InstIterator.h" |
17 | #include "llvm/IR/InstVisitor.h" |
18 | #include "llvm/IR/PassManager.h" |
19 | #include "llvm/Transforms/Utils/Cloning.h" |
20 | |
21 | #include <list> |
22 | |
23 | #define DEBUG_TYPE "spirv-regularizer" |
24 | |
25 | using namespace llvm; |
26 | |
27 | namespace { |
28 | struct SPIRVRegularizer : public FunctionPass, InstVisitor<SPIRVRegularizer> { |
29 | DenseMap<Function *, Function *> Old2NewFuncs; |
30 | |
31 | public: |
32 | static char ID; |
33 | SPIRVRegularizer() : FunctionPass(ID) {} |
34 | bool runOnFunction(Function &F) override; |
35 | StringRef getPassName() const override { return "SPIR-V Regularizer" ; } |
36 | |
37 | void getAnalysisUsage(AnalysisUsage &AU) const override { |
38 | FunctionPass::getAnalysisUsage(AU); |
39 | } |
40 | void visitCallInst(CallInst &CI); |
41 | |
42 | private: |
43 | void visitCallScalToVec(CallInst *CI, StringRef MangledName, |
44 | StringRef DemangledName); |
45 | void runLowerConstExpr(Function &F); |
46 | }; |
47 | } // namespace |
48 | |
49 | char SPIRVRegularizer::ID = 0; |
50 | |
51 | INITIALIZE_PASS(SPIRVRegularizer, DEBUG_TYPE, "SPIR-V Regularizer" , false, |
52 | false) |
53 | |
54 | // Since SPIR-V cannot represent constant expression, constant expressions |
55 | // in LLVM IR need to be lowered to instructions. For each function, |
56 | // the constant expressions used by instructions of the function are replaced |
57 | // by instructions placed in the entry block since it dominates all other BBs. |
58 | // Each constant expression only needs to be lowered once in each function |
59 | // and all uses of it by instructions in that function are replaced by |
60 | // one instruction. |
61 | // TODO: remove redundant instructions for common subexpression. |
62 | void SPIRVRegularizer::runLowerConstExpr(Function &F) { |
63 | LLVMContext &Ctx = F.getContext(); |
64 | std::list<Instruction *> WorkList; |
65 | for (auto &II : instructions(F)) |
66 | WorkList.push_back(x: &II); |
67 | |
68 | auto FBegin = F.begin(); |
69 | while (!WorkList.empty()) { |
70 | Instruction *II = WorkList.front(); |
71 | |
72 | auto LowerOp = [&II, &FBegin, &F](Value *V) -> Value * { |
73 | if (isa<Function>(Val: V)) |
74 | return V; |
75 | auto *CE = cast<ConstantExpr>(Val: V); |
76 | LLVM_DEBUG(dbgs() << "[lowerConstantExpressions] " << *CE); |
77 | auto ReplInst = CE->getAsInstruction(); |
78 | auto InsPoint = II->getParent() == &*FBegin ? II : &FBegin->back(); |
79 | ReplInst->insertBefore(InsertPos: InsPoint->getIterator()); |
80 | LLVM_DEBUG(dbgs() << " -> " << *ReplInst << '\n'); |
81 | std::vector<Instruction *> Users; |
82 | // Do not replace use during iteration of use. Do it in another loop. |
83 | for (auto U : CE->users()) { |
84 | LLVM_DEBUG(dbgs() << "[lowerConstantExpressions] Use: " << *U << '\n'); |
85 | auto InstUser = dyn_cast<Instruction>(Val: U); |
86 | // Only replace users in scope of current function. |
87 | if (InstUser && InstUser->getParent()->getParent() == &F) |
88 | Users.push_back(x: InstUser); |
89 | } |
90 | for (auto &User : Users) { |
91 | if (ReplInst->getParent() == User->getParent() && |
92 | User->comesBefore(Other: ReplInst)) |
93 | ReplInst->moveBefore(InsertPos: User->getIterator()); |
94 | User->replaceUsesOfWith(From: CE, To: ReplInst); |
95 | } |
96 | return ReplInst; |
97 | }; |
98 | |
99 | WorkList.pop_front(); |
100 | auto LowerConstantVec = [&II, &LowerOp, &WorkList, |
101 | &Ctx](ConstantVector *Vec, |
102 | unsigned NumOfOp) -> Value * { |
103 | if (std::all_of(first: Vec->op_begin(), last: Vec->op_end(), pred: [](Value *V) { |
104 | return isa<ConstantExpr>(Val: V) || isa<Function>(Val: V); |
105 | })) { |
106 | // Expand a vector of constexprs and construct it back with |
107 | // series of insertelement instructions. |
108 | std::list<Value *> OpList; |
109 | std::transform(first: Vec->op_begin(), last: Vec->op_end(), |
110 | result: std::back_inserter(x&: OpList), |
111 | unary_op: [LowerOp](Value *V) { return LowerOp(V); }); |
112 | Value *Repl = nullptr; |
113 | unsigned Idx = 0; |
114 | auto *PhiII = dyn_cast<PHINode>(Val: II); |
115 | Instruction *InsPoint = |
116 | PhiII ? &PhiII->getIncomingBlock(i: NumOfOp)->back() : II; |
117 | std::list<Instruction *> ReplList; |
118 | for (auto V : OpList) { |
119 | if (auto *Inst = dyn_cast<Instruction>(Val: V)) |
120 | ReplList.push_back(x: Inst); |
121 | Repl = InsertElementInst::Create( |
122 | Vec: (Repl ? Repl : PoisonValue::get(T: Vec->getType())), NewElt: V, |
123 | Idx: ConstantInt::get(Ty: Type::getInt32Ty(C&: Ctx), V: Idx++), NameStr: "" , |
124 | InsertBefore: InsPoint->getIterator()); |
125 | } |
126 | WorkList.splice(position: WorkList.begin(), x&: ReplList); |
127 | return Repl; |
128 | } |
129 | return nullptr; |
130 | }; |
131 | for (unsigned OI = 0, OE = II->getNumOperands(); OI != OE; ++OI) { |
132 | auto *Op = II->getOperand(i: OI); |
133 | if (auto *Vec = dyn_cast<ConstantVector>(Val: Op)) { |
134 | Value *ReplInst = LowerConstantVec(Vec, OI); |
135 | if (ReplInst) |
136 | II->replaceUsesOfWith(From: Op, To: ReplInst); |
137 | } else if (auto CE = dyn_cast<ConstantExpr>(Val: Op)) { |
138 | WorkList.push_front(x: cast<Instruction>(Val: LowerOp(CE))); |
139 | } else if (auto MDAsVal = dyn_cast<MetadataAsValue>(Val: Op)) { |
140 | auto ConstMD = dyn_cast<ConstantAsMetadata>(Val: MDAsVal->getMetadata()); |
141 | if (!ConstMD) |
142 | continue; |
143 | Constant *C = ConstMD->getValue(); |
144 | Value *ReplInst = nullptr; |
145 | if (auto *Vec = dyn_cast<ConstantVector>(Val: C)) |
146 | ReplInst = LowerConstantVec(Vec, OI); |
147 | if (auto *CE = dyn_cast<ConstantExpr>(Val: C)) |
148 | ReplInst = LowerOp(CE); |
149 | if (!ReplInst) |
150 | continue; |
151 | Metadata *RepMD = ValueAsMetadata::get(V: ReplInst); |
152 | Value *RepMDVal = MetadataAsValue::get(Context&: Ctx, MD: RepMD); |
153 | II->setOperand(i: OI, Val: RepMDVal); |
154 | WorkList.push_front(x: cast<Instruction>(Val: ReplInst)); |
155 | } |
156 | } |
157 | } |
158 | } |
159 | |
160 | // It fixes calls to OCL builtins that accept vector arguments and one of them |
161 | // is actually a scalar splat. |
162 | void SPIRVRegularizer::visitCallInst(CallInst &CI) { |
163 | auto F = CI.getCalledFunction(); |
164 | if (!F) |
165 | return; |
166 | |
167 | auto MangledName = F->getName(); |
168 | char *NameStr = itaniumDemangle(mangled_name: F->getName().data()); |
169 | if (!NameStr) |
170 | return; |
171 | StringRef DemangledName(NameStr); |
172 | |
173 | // TODO: add support for other builtins. |
174 | if (DemangledName.starts_with(Prefix: "fmin" ) || DemangledName.starts_with(Prefix: "fmax" ) || |
175 | DemangledName.starts_with(Prefix: "min" ) || DemangledName.starts_with(Prefix: "max" )) |
176 | visitCallScalToVec(CI: &CI, MangledName, DemangledName); |
177 | free(ptr: NameStr); |
178 | } |
179 | |
180 | void SPIRVRegularizer::visitCallScalToVec(CallInst *CI, StringRef MangledName, |
181 | StringRef DemangledName) { |
182 | // Check if all arguments have the same type - it's simple case. |
183 | auto Uniform = true; |
184 | Type *Arg0Ty = CI->getOperand(i_nocapture: 0)->getType(); |
185 | auto IsArg0Vector = isa<VectorType>(Val: Arg0Ty); |
186 | for (unsigned I = 1, E = CI->arg_size(); Uniform && (I != E); ++I) |
187 | Uniform = isa<VectorType>(Val: CI->getOperand(i_nocapture: I)->getType()) == IsArg0Vector; |
188 | if (Uniform) |
189 | return; |
190 | |
191 | auto *OldF = CI->getCalledFunction(); |
192 | Function *NewF = nullptr; |
193 | auto [It, Inserted] = Old2NewFuncs.try_emplace(Key: OldF); |
194 | if (Inserted) { |
195 | AttributeList Attrs = CI->getCalledFunction()->getAttributes(); |
196 | SmallVector<Type *, 2> ArgTypes = {OldF->getArg(i: 0)->getType(), Arg0Ty}; |
197 | auto *NewFTy = |
198 | FunctionType::get(Result: OldF->getReturnType(), Params: ArgTypes, isVarArg: OldF->isVarArg()); |
199 | NewF = Function::Create(Ty: NewFTy, Linkage: OldF->getLinkage(), N: OldF->getName(), |
200 | M&: *OldF->getParent()); |
201 | ValueToValueMapTy VMap; |
202 | auto NewFArgIt = NewF->arg_begin(); |
203 | for (auto &Arg : OldF->args()) { |
204 | auto ArgName = Arg.getName(); |
205 | NewFArgIt->setName(ArgName); |
206 | VMap[&Arg] = &(*NewFArgIt++); |
207 | } |
208 | SmallVector<ReturnInst *, 8> Returns; |
209 | CloneFunctionInto(NewFunc: NewF, OldFunc: OldF, VMap, |
210 | Changes: CloneFunctionChangeType::LocalChangesOnly, Returns); |
211 | NewF->setAttributes(Attrs); |
212 | It->second = NewF; |
213 | } else { |
214 | NewF = It->second; |
215 | } |
216 | assert(NewF); |
217 | |
218 | // This produces an instruction sequence that implements a splat of |
219 | // CI->getOperand(1) to a vector Arg0Ty. However, we use InsertElementInst |
220 | // and ShuffleVectorInst to generate the same code as the SPIR-V translator. |
221 | // For instance (transcoding/OpMin.ll), this call |
222 | // call spir_func <2 x i32> @_Z3minDv2_ii(<2 x i32> <i32 1, i32 10>, i32 5) |
223 | // is translated to |
224 | // %8 = OpUndef %v2uint |
225 | // %14 = OpConstantComposite %v2uint %uint_1 %uint_10 |
226 | // ... |
227 | // %10 = OpCompositeInsert %v2uint %uint_5 %8 0 |
228 | // %11 = OpVectorShuffle %v2uint %10 %8 0 0 |
229 | // %call = OpExtInst %v2uint %1 s_min %14 %11 |
230 | auto ConstInt = ConstantInt::get(Ty: IntegerType::get(C&: CI->getContext(), NumBits: 32), V: 0); |
231 | PoisonValue *PVal = PoisonValue::get(T: Arg0Ty); |
232 | Instruction *Inst = InsertElementInst::Create( |
233 | Vec: PVal, NewElt: CI->getOperand(i_nocapture: 1), Idx: ConstInt, NameStr: "" , InsertBefore: CI->getIterator()); |
234 | ElementCount VecElemCount = cast<VectorType>(Val: Arg0Ty)->getElementCount(); |
235 | Constant *ConstVec = ConstantVector::getSplat(EC: VecElemCount, Elt: ConstInt); |
236 | Value *NewVec = |
237 | new ShuffleVectorInst(Inst, PVal, ConstVec, "" , CI->getIterator()); |
238 | CI->setOperand(i_nocapture: 1, Val_nocapture: NewVec); |
239 | CI->replaceUsesOfWith(From: OldF, To: NewF); |
240 | CI->mutateFunctionType(FTy: NewF->getFunctionType()); |
241 | } |
242 | |
243 | bool SPIRVRegularizer::runOnFunction(Function &F) { |
244 | runLowerConstExpr(F); |
245 | visit(F); |
246 | for (auto &OldNew : Old2NewFuncs) { |
247 | Function *OldF = OldNew.first; |
248 | Function *NewF = OldNew.second; |
249 | NewF->takeName(V: OldF); |
250 | OldF->eraseFromParent(); |
251 | } |
252 | return true; |
253 | } |
254 | |
255 | FunctionPass *llvm::createSPIRVRegularizerPass() { |
256 | return new SPIRVRegularizer(); |
257 | } |
258 | |