1//===- DXILLegalizePass.cpp - Legalizes llvm IR for DXIL ------------------===//
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#include "DXILLegalizePass.h"
10#include "DirectX.h"
11#include "llvm/ADT/APInt.h"
12#include "llvm/IR/Constants.h"
13#include "llvm/IR/Function.h"
14#include "llvm/IR/IRBuilder.h"
15#include "llvm/IR/InstIterator.h"
16#include "llvm/IR/Instruction.h"
17#include "llvm/IR/Instructions.h"
18#include "llvm/IR/Module.h"
19#include "llvm/Pass.h"
20#include "llvm/Transforms/Utils/BasicBlockUtils.h"
21#include "llvm/Transforms/Utils/Local.h"
22#include <functional>
23
24#define DEBUG_TYPE "dxil-legalize"
25
26using namespace llvm;
27
28static bool legalizeFreeze(Instruction &I,
29 SmallVectorImpl<Instruction *> &ToRemove,
30 DenseMap<Value *, Value *>) {
31 auto *FI = dyn_cast<FreezeInst>(Val: &I);
32 if (!FI)
33 return false;
34
35 FI->replaceAllUsesWith(V: FI->getOperand(i_nocapture: 0));
36 ToRemove.push_back(Elt: FI);
37 return true;
38}
39
40static bool fixI8UseChain(Instruction &I,
41 SmallVectorImpl<Instruction *> &ToRemove,
42 DenseMap<Value *, Value *> &ReplacedValues) {
43
44 auto ProcessOperands = [&](SmallVector<Value *> &NewOperands) {
45 Type *InstrType = IntegerType::get(C&: I.getContext(), NumBits: 32);
46
47 for (unsigned OpIdx = 0; OpIdx < I.getNumOperands(); ++OpIdx) {
48 Value *Op = I.getOperand(i: OpIdx);
49 if (ReplacedValues.count(Val: Op) &&
50 ReplacedValues[Op]->getType()->isIntegerTy())
51 InstrType = ReplacedValues[Op]->getType();
52 }
53
54 for (unsigned OpIdx = 0; OpIdx < I.getNumOperands(); ++OpIdx) {
55 Value *Op = I.getOperand(i: OpIdx);
56 if (ReplacedValues.count(Val: Op))
57 NewOperands.push_back(Elt: ReplacedValues[Op]);
58 else if (auto *Imm = dyn_cast<ConstantInt>(Val: Op)) {
59 APInt Value = Imm->getValue();
60 unsigned NewBitWidth = InstrType->getIntegerBitWidth();
61 // Note: options here are sext or sextOrTrunc.
62 // Since i8 isn't supported, we assume new values
63 // will always have a higher bitness.
64 assert(NewBitWidth > Value.getBitWidth() &&
65 "Replacement's BitWidth should be larger than Current.");
66 APInt NewValue = Value.sext(width: NewBitWidth);
67 NewOperands.push_back(Elt: ConstantInt::get(Ty: InstrType, V: NewValue));
68 } else {
69 assert(!Op->getType()->isIntegerTy(8));
70 NewOperands.push_back(Elt: Op);
71 }
72 }
73 };
74 IRBuilder<> Builder(&I);
75 if (auto *Trunc = dyn_cast<TruncInst>(Val: &I)) {
76 if (Trunc->getDestTy()->isIntegerTy(BitWidth: 8)) {
77 ReplacedValues[Trunc] = Trunc->getOperand(i_nocapture: 0);
78 ToRemove.push_back(Elt: Trunc);
79 return true;
80 }
81 }
82
83 if (auto *Store = dyn_cast<StoreInst>(Val: &I)) {
84 if (!Store->getValueOperand()->getType()->isIntegerTy(BitWidth: 8))
85 return false;
86 SmallVector<Value *> NewOperands;
87 ProcessOperands(NewOperands);
88 Value *NewStore = Builder.CreateStore(Val: NewOperands[0], Ptr: NewOperands[1]);
89 ReplacedValues[Store] = NewStore;
90 ToRemove.push_back(Elt: Store);
91 return true;
92 }
93
94 if (auto *Load = dyn_cast<LoadInst>(Val: &I);
95 Load && I.getType()->isIntegerTy(BitWidth: 8)) {
96 SmallVector<Value *> NewOperands;
97 ProcessOperands(NewOperands);
98 Type *ElementType = NewOperands[0]->getType();
99 if (auto *AI = dyn_cast<AllocaInst>(Val: NewOperands[0]))
100 ElementType = AI->getAllocatedType();
101 if (auto *GEP = dyn_cast<GetElementPtrInst>(Val: NewOperands[0])) {
102 ElementType = GEP->getSourceElementType();
103 }
104 if (ElementType->isArrayTy())
105 ElementType = ElementType->getArrayElementType();
106 LoadInst *NewLoad = Builder.CreateLoad(Ty: ElementType, Ptr: NewOperands[0]);
107 ReplacedValues[Load] = NewLoad;
108 ToRemove.push_back(Elt: Load);
109 return true;
110 }
111
112 if (auto *Load = dyn_cast<LoadInst>(Val: &I);
113 Load && isa<ConstantExpr>(Val: Load->getPointerOperand())) {
114 auto *CE = dyn_cast<ConstantExpr>(Val: Load->getPointerOperand());
115 if (!(CE->getOpcode() == Instruction::GetElementPtr))
116 return false;
117 auto *GEP = dyn_cast<GEPOperator>(Val: CE);
118 if (!GEP->getSourceElementType()->isIntegerTy(BitWidth: 8))
119 return false;
120
121 Type *ElementType = Load->getType();
122 ConstantInt *Offset = dyn_cast<ConstantInt>(Val: GEP->getOperand(i_nocapture: 1));
123 uint32_t ByteOffset = Offset->getZExtValue();
124 uint32_t ElemSize = Load->getDataLayout().getTypeAllocSize(Ty: ElementType);
125 uint32_t Index = ByteOffset / ElemSize;
126
127 Value *PtrOperand = GEP->getPointerOperand();
128 Type *GEPType = GEP->getPointerOperandType();
129
130 if (auto *GV = dyn_cast<GlobalVariable>(Val: PtrOperand))
131 GEPType = GV->getValueType();
132 if (auto *AI = dyn_cast<AllocaInst>(Val: PtrOperand))
133 GEPType = AI->getAllocatedType();
134
135 if (auto *ArrTy = dyn_cast<ArrayType>(Val: GEPType))
136 GEPType = ArrTy;
137 else
138 GEPType = ArrayType::get(ElementType, NumElements: 1); // its a scalar
139
140 Value *NewGEP = Builder.CreateGEP(
141 Ty: GEPType, Ptr: PtrOperand, IdxList: {Builder.getInt32(C: 0), Builder.getInt32(C: Index)},
142 Name: GEP->getName(), NW: GEP->getNoWrapFlags());
143
144 LoadInst *NewLoad = Builder.CreateLoad(Ty: ElementType, Ptr: NewGEP);
145 ReplacedValues[Load] = NewLoad;
146 Load->replaceAllUsesWith(V: NewLoad);
147 ToRemove.push_back(Elt: Load);
148 return true;
149 }
150
151 if (auto *BO = dyn_cast<BinaryOperator>(Val: &I)) {
152 if (!I.getType()->isIntegerTy(BitWidth: 8))
153 return false;
154 SmallVector<Value *> NewOperands;
155 ProcessOperands(NewOperands);
156 Value *NewInst =
157 Builder.CreateBinOp(Opc: BO->getOpcode(), LHS: NewOperands[0], RHS: NewOperands[1]);
158 if (auto *OBO = dyn_cast<OverflowingBinaryOperator>(Val: &I)) {
159 auto *NewBO = dyn_cast<BinaryOperator>(Val: NewInst);
160 if (NewBO && OBO->hasNoSignedWrap())
161 NewBO->setHasNoSignedWrap();
162 if (NewBO && OBO->hasNoUnsignedWrap())
163 NewBO->setHasNoUnsignedWrap();
164 }
165 ReplacedValues[BO] = NewInst;
166 ToRemove.push_back(Elt: BO);
167 return true;
168 }
169
170 if (auto *Sel = dyn_cast<SelectInst>(Val: &I)) {
171 if (!I.getType()->isIntegerTy(BitWidth: 8))
172 return false;
173 SmallVector<Value *> NewOperands;
174 ProcessOperands(NewOperands);
175 Value *NewInst = Builder.CreateSelect(C: Sel->getCondition(), True: NewOperands[1],
176 False: NewOperands[2]);
177 ReplacedValues[Sel] = NewInst;
178 ToRemove.push_back(Elt: Sel);
179 return true;
180 }
181
182 if (auto *Cmp = dyn_cast<CmpInst>(Val: &I)) {
183 if (!Cmp->getOperand(i_nocapture: 0)->getType()->isIntegerTy(BitWidth: 8))
184 return false;
185 SmallVector<Value *> NewOperands;
186 ProcessOperands(NewOperands);
187 Value *NewInst =
188 Builder.CreateCmp(Pred: Cmp->getPredicate(), LHS: NewOperands[0], RHS: NewOperands[1]);
189 Cmp->replaceAllUsesWith(V: NewInst);
190 ReplacedValues[Cmp] = NewInst;
191 ToRemove.push_back(Elt: Cmp);
192 return true;
193 }
194
195 if (auto *Cast = dyn_cast<CastInst>(Val: &I)) {
196 if (!Cast->getSrcTy()->isIntegerTy(BitWidth: 8))
197 return false;
198
199 ToRemove.push_back(Elt: Cast);
200 auto *Replacement = ReplacedValues[Cast->getOperand(i_nocapture: 0)];
201 if (Cast->getType() == Replacement->getType()) {
202 Cast->replaceAllUsesWith(V: Replacement);
203 return true;
204 }
205
206 Value *AdjustedCast = nullptr;
207 if (Cast->getOpcode() == Instruction::ZExt)
208 AdjustedCast = Builder.CreateZExtOrTrunc(V: Replacement, DestTy: Cast->getType());
209 if (Cast->getOpcode() == Instruction::SExt)
210 AdjustedCast = Builder.CreateSExtOrTrunc(V: Replacement, DestTy: Cast->getType());
211
212 if (AdjustedCast)
213 Cast->replaceAllUsesWith(V: AdjustedCast);
214 }
215 if (auto *GEP = dyn_cast<GetElementPtrInst>(Val: &I)) {
216 if (!GEP->getType()->isPointerTy() ||
217 !GEP->getSourceElementType()->isIntegerTy(BitWidth: 8))
218 return false;
219
220 Value *BasePtr = GEP->getPointerOperand();
221 if (ReplacedValues.count(Val: BasePtr))
222 BasePtr = ReplacedValues[BasePtr];
223
224 Type *ElementType = BasePtr->getType();
225
226 if (auto *AI = dyn_cast<AllocaInst>(Val: BasePtr))
227 ElementType = AI->getAllocatedType();
228 if (auto *GV = dyn_cast<GlobalVariable>(Val: BasePtr))
229 ElementType = GV->getValueType();
230
231 Type *GEPType = ElementType;
232 if (auto *ArrTy = dyn_cast<ArrayType>(Val: ElementType))
233 ElementType = ArrTy->getArrayElementType();
234 else
235 GEPType = ArrayType::get(ElementType, NumElements: 1); // its a scalar
236
237 ConstantInt *Offset = dyn_cast<ConstantInt>(Val: GEP->getOperand(i_nocapture: 1));
238 // Note: i8 to i32 offset conversion without emitting IR requires constant
239 // ints. Since offset conversion is common, we can safely assume Offset is
240 // always a ConstantInt, so no need to have a conditional bail out on
241 // nullptr, instead assert this is the case.
242 assert(Offset && "Offset is expected to be a ConstantInt");
243 uint32_t ByteOffset = Offset->getZExtValue();
244 uint32_t ElemSize = GEP->getDataLayout().getTypeAllocSize(Ty: ElementType);
245 assert(ElemSize > 0 && "ElementSize must be set");
246 uint32_t Index = ByteOffset / ElemSize;
247 Value *NewGEP = Builder.CreateGEP(
248 Ty: GEPType, Ptr: BasePtr, IdxList: {Builder.getInt32(C: 0), Builder.getInt32(C: Index)},
249 Name: GEP->getName(), NW: GEP->getNoWrapFlags());
250 ReplacedValues[GEP] = NewGEP;
251 GEP->replaceAllUsesWith(V: NewGEP);
252 ToRemove.push_back(Elt: GEP);
253 return true;
254 }
255 return false;
256}
257
258static bool upcastI8AllocasAndUses(Instruction &I,
259 SmallVectorImpl<Instruction *> &ToRemove,
260 DenseMap<Value *, Value *> &ReplacedValues) {
261 auto *AI = dyn_cast<AllocaInst>(Val: &I);
262 if (!AI || !AI->getAllocatedType()->isIntegerTy(BitWidth: 8))
263 return false;
264
265 Type *SmallestType = nullptr;
266
267 auto ProcessLoad = [&](LoadInst *Load) {
268 for (User *LU : Load->users()) {
269 CastInst *Cast = dyn_cast<CastInst>(Val: LU);
270 if (!Cast)
271 continue;
272 Type *Ty = Cast->getType();
273
274 if (!SmallestType ||
275 Ty->getPrimitiveSizeInBits() < SmallestType->getPrimitiveSizeInBits())
276 SmallestType = Ty;
277 }
278 };
279
280 for (User *U : AI->users()) {
281 if (auto *Load = dyn_cast<LoadInst>(Val: U))
282 ProcessLoad(Load);
283 else if (auto *GEP = dyn_cast<GetElementPtrInst>(Val: U)) {
284 for (User *GU : GEP->users()) {
285 if (auto *Load = dyn_cast<LoadInst>(Val: GU))
286 ProcessLoad(Load);
287 }
288 }
289 }
290
291 if (!SmallestType)
292 return false; // no valid casts found
293
294 // Replace alloca
295 IRBuilder<> Builder(AI);
296 auto *NewAlloca = Builder.CreateAlloca(Ty: SmallestType);
297 ReplacedValues[AI] = NewAlloca;
298 ToRemove.push_back(Elt: AI);
299 return true;
300}
301
302static bool
303downcastI64toI32InsertExtractElements(Instruction &I,
304 SmallVectorImpl<Instruction *> &ToRemove,
305 DenseMap<Value *, Value *> &) {
306
307 if (auto *Extract = dyn_cast<ExtractElementInst>(Val: &I)) {
308 Value *Idx = Extract->getIndexOperand();
309 auto *CI = dyn_cast<ConstantInt>(Val: Idx);
310 if (CI && CI->getBitWidth() == 64) {
311 IRBuilder<> Builder(Extract);
312 int64_t IndexValue = CI->getSExtValue();
313 auto *Idx32 =
314 ConstantInt::get(Ty: Type::getInt32Ty(C&: I.getContext()), V: IndexValue);
315 Value *NewExtract = Builder.CreateExtractElement(
316 Vec: Extract->getVectorOperand(), Idx: Idx32, Name: Extract->getName());
317
318 Extract->replaceAllUsesWith(V: NewExtract);
319 ToRemove.push_back(Elt: Extract);
320 return true;
321 }
322 }
323
324 if (auto *Insert = dyn_cast<InsertElementInst>(Val: &I)) {
325 Value *Idx = Insert->getOperand(i_nocapture: 2);
326 auto *CI = dyn_cast<ConstantInt>(Val: Idx);
327 if (CI && CI->getBitWidth() == 64) {
328 int64_t IndexValue = CI->getSExtValue();
329 auto *Idx32 =
330 ConstantInt::get(Ty: Type::getInt32Ty(C&: I.getContext()), V: IndexValue);
331 IRBuilder<> Builder(Insert);
332 Value *Insert32Index = Builder.CreateInsertElement(
333 Vec: Insert->getOperand(i_nocapture: 0), NewElt: Insert->getOperand(i_nocapture: 1), Idx: Idx32,
334 Name: Insert->getName());
335
336 Insert->replaceAllUsesWith(V: Insert32Index);
337 ToRemove.push_back(Elt: Insert);
338 return true;
339 }
340 }
341 return false;
342}
343
344static bool updateFnegToFsub(Instruction &I,
345 SmallVectorImpl<Instruction *> &ToRemove,
346 DenseMap<Value *, Value *> &) {
347 const Intrinsic::ID ID = I.getOpcode();
348 if (ID != Instruction::FNeg)
349 return false;
350
351 IRBuilder<> Builder(&I);
352 Value *In = I.getOperand(i: 0);
353 Value *Zero = ConstantFP::get(Ty: In->getType(), V: -0.0);
354 I.replaceAllUsesWith(V: Builder.CreateFSub(L: Zero, R: In));
355 ToRemove.push_back(Elt: &I);
356 return true;
357}
358
359static bool
360resolveUnreachableSwitchDefault(Instruction &I,
361 SmallVectorImpl<Instruction *> &ToRemove,
362 DenseMap<Value *, Value *> &) {
363 auto *SI = dyn_cast<SwitchInst>(Val: &I);
364 if (!SI || SI->getNumCases() == 0)
365 return false;
366
367 BasicBlock *DefaultBB = SI->getDefaultDest();
368
369 // Check if the default destination ends with an unreachable instruction.
370 if (DefaultBB->size() == 0 ||
371 !isa<UnreachableInst>(Val: DefaultBB->getTerminator()))
372 return false;
373
374 // Try to find a common successor of all case destinations. If all case
375 // blocks unconditionally branch to the same block, that is the common
376 // successor. This is just a best effort, and is done as the original form of
377 // the switch statement was likely in this form before being transformed to
378 // an unreachable branch.
379 BasicBlock *CommonSuccessor = nullptr;
380 for (auto &Case : SI->cases()) {
381 BasicBlock *CaseBB = Case.getCaseSuccessor();
382 auto *BI = dyn_cast<UncondBrInst>(Val: CaseBB->getTerminator());
383 if (!BI) {
384 CommonSuccessor = nullptr;
385 break;
386 }
387 BasicBlock *Succ = BI->getSuccessor(i: 0);
388 if (!CommonSuccessor)
389 CommonSuccessor = Succ;
390 else if (CommonSuccessor != Succ) {
391 CommonSuccessor = nullptr;
392 break;
393 }
394 }
395
396 BasicBlock *NewDefault =
397 CommonSuccessor ? CommonSuccessor : SI->case_begin()->getCaseSuccessor();
398
399 BasicBlock *SwitchBB = SI->getParent();
400 SI->setDefaultDest(NewDefault);
401
402 // Ensure all phi nodes are legal by adding an incoming poison value from the
403 // unreachable branch.
404 for (PHINode &Phi : NewDefault->phis())
405 Phi.addIncoming(V: PoisonValue::get(T: Phi.getType()), BB: SwitchBB);
406
407 return true;
408}
409
410static bool
411legalizeScalarLoadStoreOnArrays(Instruction &I,
412 SmallVectorImpl<Instruction *> &ToRemove,
413 DenseMap<Value *, Value *> &) {
414
415 Value *PtrOp;
416 unsigned PtrOpIndex;
417 [[maybe_unused]] Type *LoadStoreTy;
418 if (auto *LI = dyn_cast<LoadInst>(Val: &I)) {
419 PtrOp = LI->getPointerOperand();
420 PtrOpIndex = LI->getPointerOperandIndex();
421 LoadStoreTy = LI->getType();
422 } else if (auto *SI = dyn_cast<StoreInst>(Val: &I)) {
423 PtrOp = SI->getPointerOperand();
424 PtrOpIndex = SI->getPointerOperandIndex();
425 LoadStoreTy = SI->getValueOperand()->getType();
426 } else
427 return false;
428
429 // If the load/store is not of a single-value type (i.e., scalar or vector)
430 // then we do not modify it. It shouldn't be a vector either because the
431 // dxil-data-scalarization pass is expected to run before this, but it's not
432 // incorrect to apply this transformation to vector load/stores.
433 if (!LoadStoreTy->isSingleValueType())
434 return false;
435
436 Type *ArrayTy;
437 if (auto *GlobalVarPtrOp = dyn_cast<GlobalVariable>(Val: PtrOp))
438 ArrayTy = GlobalVarPtrOp->getValueType();
439 else if (auto *AllocaPtrOp = dyn_cast<AllocaInst>(Val: PtrOp))
440 ArrayTy = AllocaPtrOp->getAllocatedType();
441 else
442 return false;
443
444 if (!isa<ArrayType>(Val: ArrayTy))
445 return false;
446
447 assert(ArrayTy->getArrayElementType() == LoadStoreTy &&
448 "Expected array element type to be the same as to the scalar load or "
449 "store type");
450
451 Value *Zero = ConstantInt::get(Ty: Type::getInt32Ty(C&: I.getContext()), V: 0);
452 Value *GEP = GetElementPtrInst::Create(
453 PointeeType: ArrayTy, Ptr: PtrOp, IdxList: {Zero, Zero}, NW: GEPNoWrapFlags::all(), NameStr: "", InsertBefore: I.getIterator());
454 I.setOperand(i: PtrOpIndex, Val: GEP);
455 return true;
456}
457
458namespace {
459class DXILLegalizationPipeline {
460
461public:
462 DXILLegalizationPipeline() { initializeLegalizationPipeline(); }
463
464 bool runLegalizationPipeline(Function &F) {
465 bool MadeChange = false;
466 SmallVector<Instruction *> ToRemove;
467 DenseMap<Value *, Value *> ReplacedValues;
468 for (int Stage = 0; Stage < NumStages; ++Stage) {
469 ToRemove.clear();
470 ReplacedValues.clear();
471 for (auto &I : instructions(F)) {
472 for (auto &LegalizationFn : LegalizationPipeline[Stage])
473 MadeChange |= LegalizationFn(I, ToRemove, ReplacedValues);
474 }
475
476 for (auto *Inst : reverse(C&: ToRemove))
477 Inst->eraseFromParent();
478 }
479
480 if (MadeChange)
481 MadeChange |= removeUnreachableBlocks(F);
482 return MadeChange;
483 }
484
485private:
486 enum LegalizationStage { Stage1 = 0, Stage2 = 1, NumStages };
487
488 using LegalizationFnTy =
489 std::function<bool(Instruction &, SmallVectorImpl<Instruction *> &,
490 DenseMap<Value *, Value *> &)>;
491
492 SmallVector<LegalizationFnTy> LegalizationPipeline[NumStages];
493
494 void initializeLegalizationPipeline() {
495 LegalizationPipeline[Stage1].push_back(Elt: upcastI8AllocasAndUses);
496 LegalizationPipeline[Stage1].push_back(Elt: fixI8UseChain);
497 LegalizationPipeline[Stage1].push_back(Elt: legalizeFreeze);
498 LegalizationPipeline[Stage1].push_back(Elt: updateFnegToFsub);
499 LegalizationPipeline[Stage1].push_back(
500 Elt: downcastI64toI32InsertExtractElements);
501 LegalizationPipeline[Stage2].push_back(Elt: legalizeScalarLoadStoreOnArrays);
502 LegalizationPipeline[Stage2].push_back(Elt: resolveUnreachableSwitchDefault);
503 }
504};
505
506class DXILLegalizeLegacy : public FunctionPass {
507
508public:
509 bool runOnFunction(Function &F) override;
510 DXILLegalizeLegacy() : FunctionPass(ID) {}
511
512 static char ID; // Pass identification.
513};
514} // namespace
515
516PreservedAnalyses DXILLegalizePass::run(Function &F,
517 FunctionAnalysisManager &FAM) {
518 DXILLegalizationPipeline DXLegalize;
519 bool MadeChanges = DXLegalize.runLegalizationPipeline(F);
520 if (!MadeChanges)
521 return PreservedAnalyses::all();
522 PreservedAnalyses PA;
523 return PA;
524}
525
526bool DXILLegalizeLegacy::runOnFunction(Function &F) {
527 DXILLegalizationPipeline DXLegalize;
528 return DXLegalize.runLegalizationPipeline(F);
529}
530
531char DXILLegalizeLegacy::ID = 0;
532
533INITIALIZE_PASS_BEGIN(DXILLegalizeLegacy, DEBUG_TYPE, "DXIL Legalizer", false,
534 false)
535INITIALIZE_PASS_END(DXILLegalizeLegacy, DEBUG_TYPE, "DXIL Legalizer", false,
536 false)
537
538FunctionPass *llvm::createDXILLegalizeLegacyPass() {
539 return new DXILLegalizeLegacy();
540}
541