1//===- LoadStoreVec.cpp - Vectorizer pass short load-store chains ---------===//
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 "llvm/Transforms/Vectorize/SandboxVectorizer/Passes/LoadStoreVec.h"
10#include "llvm/ADT/DenseSet.h"
11#include "llvm/SandboxIR/Instruction.h"
12#include "llvm/SandboxIR/Module.h"
13#include "llvm/SandboxIR/Region.h"
14#include "llvm/Support/CommandLine.h"
15#include "llvm/Support/InstructionCost.h"
16#include "llvm/Transforms/Vectorize/SandboxVectorizer/Debug.h"
17#include "llvm/Transforms/Vectorize/SandboxVectorizer/Legality.h"
18#include "llvm/Transforms/Vectorize/SandboxVectorizer/RegionWithScore.h"
19#include "llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h"
20#include "llvm/Transforms/Vectorize/SandboxVectorizer/VecUtils.h"
21
22namespace llvm {
23
24extern cl::opt<int> CostThreshold; // Defined in TransactionAcceptOrRevert.cpp
25
26namespace sandboxir {
27
28#define DEBUG_PREFIX_LOCAL DEBUG_PREFIX "LoadStoreVec: "
29
30std::optional<Type *> LoadStoreVec::canVectorize(BndlRef<Instruction *> Bndl) {
31 // Check if in the same BB.
32 if (LegalityAnalysis::differentBlock(Instrs: Bndl))
33 return std::nullopt;
34
35 // Check if instructions repeat.
36 if (!LegalityAnalysis::areUnique(Values: Bndl))
37 return std::nullopt;
38
39 // Check scheduling.
40 if (!Sched->trySchedule(Instrs: Bndl))
41 return std::nullopt;
42
43 return VecUtils::getCombinedVectorTypeFor(Bndl, DL: *DL);
44}
45
46void LoadStoreVec::saveIR(Region &R) {
47 Rgn = &R;
48 const auto &SB = cast<RegionWithScore>(Val: Rgn)->getScoreboard();
49 CostBefore = SB.getAfterCost() - SB.getBeforeCost();
50 Rgn->getContext().save();
51}
52
53bool LoadStoreVec::acceptOrRevert() {
54 const auto &SB = cast<RegionWithScore>(Val&: *Rgn).getScoreboard();
55 InstructionCost CostAfter = SB.getAfterCost() - SB.getBeforeCost();
56 InstructionCost CostGain = CostAfter - CostBefore;
57 LLVM_DEBUG(dbgs() << DEBUG_PREFIX_LOCAL << "CostGain=" << CostGain
58 << " (After=" << CostAfter << " Before=" << CostBefore
59 << ")\n");
60 if (CostGain > CostThreshold) {
61 LLVM_DEBUG(dbgs() << DEBUG_PREFIX_LOCAL << "Not profitable, reverting.\n");
62 Ctx->revert();
63 return false;
64 }
65 LLVM_DEBUG(dbgs() << DEBUG_PREFIX_LOCAL << "Profitable accepting.\n");
66 Ctx->accept();
67 return true;
68}
69
70LoadInst *LoadStoreVec::createVectorLoad(BndlRef<Instruction *> Loads) {
71 if (!VecUtils::areConsecutive<LoadInst, Instruction>(
72 Bndl: Loads, SE&: A->getScalarEvolution(), DL: *DL))
73 return nullptr;
74 if (!canVectorize(Bndl: Loads))
75 return nullptr;
76
77 Type *Ty = VecUtils::getCombinedVectorTypeFor(Bndl: Loads, DL: *DL);
78 Value *LdPtr = cast<LoadInst>(Val: Loads[0])->getPointerOperand();
79 // TODO: Compute alignment.
80 Align LdAlign(1);
81 auto LdWhereIt = std::next(x: VecUtils::getLowest(Instrs: Loads)->getIterator());
82 return LoadInst::create(Ty, Ptr: LdPtr, Align: LdAlign, Pos: LdWhereIt, Ctx&: *Ctx, Name: "VecIinitL");
83}
84
85Constant *LoadStoreVec::createConstantVector(BndlRef<Constant *> Operands) {
86 SmallVector<Constant *, 8> Constants;
87 Constants.reserve(N: Operands.size());
88 for (Constant *COp : Operands) {
89 if (auto *AggrCOp = dyn_cast<ConstantAggregate>(Val: COp)) {
90 // If the operand is a constant aggregate, then append all its elements.
91 for (Value *Elm : AggrCOp->operands())
92 Constants.push_back(Elt: cast<Constant>(Val: Elm));
93 } else if (auto *SeqCOp = dyn_cast<ConstantDataSequential>(Val: COp)) {
94 for (auto ElmIdx : seq<unsigned>(Size: SeqCOp->getNumElements()))
95 Constants.push_back(Elt: SeqCOp->getElementAsConstant(ElmIdx));
96 } else if (auto *Zero = dyn_cast<ConstantAggregateZero>(Val: COp)) {
97 auto *ZeroElm = Zero->getSequentialElement();
98 for ([[maybe_unused]] auto Cnt :
99 seq<unsigned>(Size: Zero->getElementCount().getFixedValue()))
100 Constants.push_back(Elt: ZeroElm);
101 } else if (isa<ConstantInt>(Val: COp) && isa<VectorType>(Val: COp->getType())) {
102 auto *Elm = ConstantInt::get(Ctx&: *Ctx, V: cast<ConstantInt>(Val: COp)->getValue());
103 for ([[maybe_unused]] auto Cnt :
104 seq<unsigned>(Size: cast<VectorType>(Val: COp->getType())
105 ->getElementCount()
106 .getFixedValue()))
107 Constants.push_back(Elt: Elm);
108 } else if (isa<ConstantFP>(Val: COp) && isa<VectorType>(Val: COp->getType())) {
109 auto *Elm = ConstantFP::get(V: cast<ConstantFP>(Val: COp)->getValue(), Ctx&: *Ctx);
110 for ([[maybe_unused]] auto Cnt :
111 seq<unsigned>(Size: cast<VectorType>(Val: COp->getType())
112 ->getElementCount()
113 .getFixedValue()))
114 Constants.push_back(Elt: Elm);
115 } else {
116 Constants.push_back(Elt: COp);
117 }
118 }
119 return ConstantVector::get(V: Constants);
120}
121
122bool LoadStoreVec::vectorizeStores(BndlRef<Instruction *> Stores, Region &Rgn) {
123 if (!VecUtils::areConsecutive<StoreInst, Instruction>(
124 Bndl: Stores, SE&: A->getScalarEvolution(), DL: *DL))
125 return false;
126 if (!canVectorize(Bndl: Stores))
127 return false;
128 SmallVector<Value *, 4> Operands;
129 Operands.reserve(N: Stores.size());
130 for (auto *I : Stores) {
131 auto *Op = cast<StoreInst>(Val: I)->getValueOperand();
132 Operands.push_back(Elt: Op);
133 }
134 BasicBlock *BB = Stores[0]->getParent();
135 // TODO: For now we only support load operands.
136 // TODO: For now we don't cross BBs.
137 // TODO: For now don't vectorize if the loads have external uses.
138 bool AllLoads = all_of(Range&: Operands, P: [BB](Value *V) {
139 auto *LI = dyn_cast<LoadInst>(Val: V);
140 if (LI == nullptr)
141 return false;
142 // TODO: For now we don't cross BBs.
143 if (LI->getParent() != BB)
144 return false;
145 if (LI->hasNUsesOrMore(Num: 2))
146 return false;
147 return true;
148 });
149 bool AllConstants =
150 all_of(Range&: Operands, P: [](Value *V) { return isa<Constant>(Val: V); });
151 if (!AllLoads && !AllConstants)
152 return false;
153
154 // Vectorizing mixed floats and integers with external uses may not be
155 // profitable on some targets, so save state here.
156 saveIR(R&: Rgn);
157 Value *VecOp = nullptr;
158 SmallVector<Instruction *, 8> Loads;
159 if (AllLoads) {
160 Loads.reserve(N: Operands.size());
161 for (Value *Op : Operands)
162 Loads.push_back(Elt: cast<Instruction>(Val: Op));
163 VecOp = createVectorLoad(Loads);
164 if (VecOp == nullptr) {
165 Ctx->accept();
166 return false;
167 }
168 } else if (AllConstants) {
169 SmallVector<Constant *, 8> Constants;
170 Constants.reserve(N: Operands.size());
171 for (Value *Op : Operands)
172 Constants.push_back(Elt: cast<Constant>(Val: Op));
173 VecOp = createConstantVector(Operands: Constants);
174 }
175
176 // Generate vector store.
177 Value *StPtr = cast<StoreInst>(Val: Stores[0])->getPointerOperand();
178 // TODO: Compute alignment.
179 Align StAlign(1);
180 auto StWhereIt = std::next(x: VecUtils::getLowest(Instrs: Stores)->getIterator());
181 StoreInst::create(V: VecOp, Ptr: StPtr, Align: StAlign, Pos: StWhereIt, Ctx&: *Ctx);
182
183 DeadInstrMorgue.collectPotentiallyDeadInstrs(Bndl: Stores);
184 if (AllLoads)
185 DeadInstrMorgue.collectPotentiallyDeadInstrs<Instruction>(Bndl: Loads);
186 DeadInstrMorgue.tryEraseDeadInstrs();
187
188 return acceptOrRevert();
189}
190
191LoadInst *LoadStoreVec::vectorizeLoads(BndlRef<Instruction *> Loads,
192 Region &Rgn) {
193 if (!VecUtils::areConsecutive<LoadInst, Instruction>(
194 Bndl: Loads, SE&: A->getScalarEvolution(), DL: *DL))
195 return nullptr;
196 auto VecTy = canVectorize(Bndl: Loads);
197 if (!VecTy)
198 return nullptr;
199
200 // TODO: Support mixed-type top-level load chains.
201 Type *VecElemTy = cast<FixedVectorType>(Val: *VecTy)->getElementType();
202 if (!all_of(Range&: Loads, P: [VecElemTy](Instruction *I) {
203 return VecUtils::getElementType(Ty: I->getType()) == VecElemTy;
204 }))
205 return nullptr;
206
207 saveIR(R&: Rgn);
208
209 auto *VecLoad = createVectorLoad(Loads);
210 if (VecLoad == nullptr) {
211 Ctx->accept();
212 return nullptr;
213 }
214
215 BasicBlock::iterator WhereIt = std::next(x: VecLoad->getIterator());
216 for (auto [Lane, OrigV] : VecUtils::enumerateLanes(Range: Loads)) {
217 auto *OrigLoad = cast<LoadInst>(Val: OrigV);
218 if (OrigLoad->hasNUses(Num: 0))
219 continue;
220 Value *Unpacked =
221 VecUtils::unpack(FromVec: VecLoad, ExtrTy: OrigLoad->getType(), Lane, WhereIt);
222 OrigLoad->replaceAllUsesWith(Other: Unpacked);
223 }
224
225 DeadInstrMorgue.collectPotentiallyDeadInstrs(Bndl: Loads);
226 DeadInstrMorgue.tryEraseDeadInstrs();
227
228 if (!acceptOrRevert())
229 return nullptr;
230 return VecLoad;
231}
232
233bool LoadStoreVec::runOnRegion(Region &Rgn, const Analyses &RegionAnalyses) {
234 SmallVector<Instruction *, 8> Bndl(Rgn.getAux().begin(), Rgn.getAux().end());
235 if (Bndl.size() < 2)
236 return false;
237 Function &F = *Bndl[0]->getParent()->getParent();
238 DL = &F.getParent()->getDataLayout();
239 Ctx = &F.getContext();
240 A = &RegionAnalyses;
241 Sched =
242 std::make_unique<Scheduler>(args&: A->getAA(), args&: *Ctx, args: SchedDirection::BottomUp);
243
244 auto Opc = Bndl[0]->getOpcode();
245 assert(
246 all_of(Bndl, [Opc](Instruction *I) { return I->getOpcode() == Opc; }) &&
247 "Expected a homogeneous seed slice!");
248
249 bool Changed = false;
250 switch (Opc) {
251 case Instruction::Opcode::Load:
252 Changed = vectorizeLoads(Loads: Bndl, Rgn) != nullptr;
253 break;
254 case Instruction::Opcode::Store:
255 Changed = vectorizeStores(Stores: Bndl, Rgn);
256 break;
257 default:
258 llvm_unreachable("Expected Load or Store");
259 }
260 Sched.reset();
261 return Changed;
262}
263
264} // namespace sandboxir
265
266} // namespace llvm
267