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
85Value *LoadStoreVec::createConstantVector(BndlRef<Value *> Operands) {
86 SmallVector<Constant *, 8> Constants;
87 Constants.reserve(N: Operands.size());
88 for (Value *Op : Operands) {
89 auto *COp = cast<Constant>(Val: Op);
90 if (auto *AggrCOp = dyn_cast<ConstantAggregate>(Val: COp)) {
91 // If the operand is a constant aggregate, then append all its elements.
92 for (Value *Elm : AggrCOp->operands())
93 Constants.push_back(Elt: cast<Constant>(Val: Elm));
94 } else if (auto *SeqCOp = dyn_cast<ConstantDataSequential>(Val: COp)) {
95 for (auto ElmIdx : seq<unsigned>(Size: SeqCOp->getNumElements()))
96 Constants.push_back(Elt: SeqCOp->getElementAsConstant(ElmIdx));
97 } else if (auto *Zero = dyn_cast<ConstantAggregateZero>(Val: COp)) {
98 auto *ZeroElm = Zero->getSequentialElement();
99 for ([[maybe_unused]] auto Cnt :
100 seq<unsigned>(Size: Zero->getElementCount().getFixedValue()))
101 Constants.push_back(Elt: ZeroElm);
102 } else if (isa<ConstantInt>(Val: COp) && isa<VectorType>(Val: COp->getType())) {
103 auto *Elm = ConstantInt::get(Ctx&: *Ctx, V: cast<ConstantInt>(Val: COp)->getValue());
104 for ([[maybe_unused]] auto Cnt :
105 seq<unsigned>(Size: cast<VectorType>(Val: COp->getType())
106 ->getElementCount()
107 .getFixedValue()))
108 Constants.push_back(Elt: Elm);
109 } else if (isa<ConstantFP>(Val: COp) && isa<VectorType>(Val: COp->getType())) {
110 auto *Elm = ConstantFP::get(V: cast<ConstantFP>(Val: COp)->getValue(), Ctx&: *Ctx);
111 for ([[maybe_unused]] auto Cnt :
112 seq<unsigned>(Size: cast<VectorType>(Val: COp->getType())
113 ->getElementCount()
114 .getFixedValue()))
115 Constants.push_back(Elt: Elm);
116 } else {
117 Constants.push_back(Elt: COp);
118 }
119 }
120 return ConstantVector::get(V: Constants);
121}
122
123bool LoadStoreVec::vectorizeStores(BndlRef<Instruction *> Stores, Region &Rgn) {
124 if (!VecUtils::areConsecutive<StoreInst, Instruction>(
125 Bndl: Stores, SE&: A->getScalarEvolution(), DL: *DL))
126 return false;
127 if (!canVectorize(Bndl: Stores))
128 return false;
129 SmallVector<Value *, 4> Operands;
130 Operands.reserve(N: Stores.size());
131 for (auto *I : Stores) {
132 auto *Op = cast<StoreInst>(Val: I)->getValueOperand();
133 Operands.push_back(Elt: Op);
134 }
135 BasicBlock *BB = Stores[0]->getParent();
136 // TODO: For now we only support load operands.
137 // TODO: For now we don't cross BBs.
138 // TODO: For now don't vectorize if the loads have external uses.
139 bool AllLoads = all_of(Range&: Operands, P: [BB](Value *V) {
140 auto *LI = dyn_cast<LoadInst>(Val: V);
141 if (LI == nullptr)
142 return false;
143 // TODO: For now we don't cross BBs.
144 if (LI->getParent() != BB)
145 return false;
146 if (LI->hasNUsesOrMore(Num: 2))
147 return false;
148 return true;
149 });
150 bool AllConstants =
151 all_of(Range&: Operands, P: [](Value *V) { return isa<Constant>(Val: V); });
152 if (!AllLoads && !AllConstants)
153 return false;
154
155 // Vectorizing mixed floats and integers with external uses may not be
156 // profitable on some targets, so save state here.
157 saveIR(R&: Rgn);
158 Value *VecOp = nullptr;
159 if (AllLoads) {
160 // TODO: Try to avoid the extra copy to an instruction vector.
161 SmallVector<Instruction *, 8> Loads;
162 Loads.reserve(N: Operands.size());
163 for (Value *Op : Operands)
164 Loads.push_back(Elt: cast<Instruction>(Val: Op));
165 VecOp = createVectorLoad(Loads);
166 if (VecOp == nullptr) {
167 Ctx->accept();
168 return false;
169 }
170 } else if (AllConstants) {
171 VecOp = createConstantVector(Operands);
172 }
173
174 // Generate vector store.
175 Value *StPtr = cast<StoreInst>(Val: Stores[0])->getPointerOperand();
176 // TODO: Compute alignment.
177 Align StAlign(1);
178 auto StWhereIt = std::next(x: VecUtils::getLowest(Instrs: Stores)->getIterator());
179 StoreInst::create(V: VecOp, Ptr: StPtr, Align: StAlign, Pos: StWhereIt, Ctx&: *Ctx);
180
181 DeadInstrMorgue.collectPotentiallyDeadInstrs(Bndl: Stores);
182 if (AllLoads)
183 DeadInstrMorgue.collectPotentiallyDeadInstrs<Value>(Bndl: Operands);
184 DeadInstrMorgue.tryEraseDeadInstrs();
185
186 return acceptOrRevert();
187}
188
189LoadInst *LoadStoreVec::vectorizeLoads(BndlRef<Instruction *> Loads,
190 Region &Rgn) {
191 if (!VecUtils::areConsecutive<LoadInst, Instruction>(
192 Bndl: Loads, SE&: A->getScalarEvolution(), DL: *DL))
193 return nullptr;
194 auto VecTy = canVectorize(Bndl: Loads);
195 if (!VecTy)
196 return nullptr;
197
198 // TODO: Support mixed-type top-level load chains.
199 Type *VecElemTy = cast<FixedVectorType>(Val: *VecTy)->getElementType();
200 if (!all_of(Range&: Loads, P: [VecElemTy](Instruction *I) {
201 return VecUtils::getElementType(Ty: I->getType()) == VecElemTy;
202 }))
203 return nullptr;
204
205 saveIR(R&: Rgn);
206
207 auto *VecLoad = createVectorLoad(Loads);
208 if (VecLoad == nullptr) {
209 Ctx->accept();
210 return nullptr;
211 }
212
213 BasicBlock::iterator WhereIt = std::next(x: VecLoad->getIterator());
214 for (auto [Lane, OrigV] : VecUtils::enumerateLanes(Range: Loads)) {
215 auto *OrigLoad = cast<LoadInst>(Val: OrigV);
216 if (OrigLoad->hasNUses(Num: 0))
217 continue;
218 Value *Unpacked =
219 VecUtils::unpack(FromVec: VecLoad, ExtrTy: OrigLoad->getType(), Lane, WhereIt);
220 OrigLoad->replaceAllUsesWith(Other: Unpacked);
221 }
222
223 DeadInstrMorgue.collectPotentiallyDeadInstrs(Bndl: Loads);
224 DeadInstrMorgue.tryEraseDeadInstrs();
225
226 if (!acceptOrRevert())
227 return nullptr;
228 return VecLoad;
229}
230
231bool LoadStoreVec::runOnRegion(Region &Rgn, const Analyses &RegionAnalyses) {
232 SmallVector<Instruction *, 8> Bndl(Rgn.getAux().begin(), Rgn.getAux().end());
233 if (Bndl.size() < 2)
234 return false;
235 Function &F = *Bndl[0]->getParent()->getParent();
236 DL = &F.getParent()->getDataLayout();
237 Ctx = &F.getContext();
238 A = &RegionAnalyses;
239 Sched =
240 std::make_unique<Scheduler>(args&: A->getAA(), args&: *Ctx, args: SchedDirection::BottomUp);
241
242 auto Opc = Bndl[0]->getOpcode();
243 assert(
244 all_of(Bndl, [Opc](Instruction *I) { return I->getOpcode() == Opc; }) &&
245 "Expected a homogeneous seed slice!");
246
247 bool Changed = false;
248 switch (Opc) {
249 case Instruction::Opcode::Load:
250 Changed = vectorizeLoads(Loads: Bndl, Rgn) != nullptr;
251 break;
252 case Instruction::Opcode::Store:
253 Changed = vectorizeStores(Stores: Bndl, Rgn);
254 break;
255 default:
256 llvm_unreachable("Expected Load or Store");
257 }
258 Sched.reset();
259 return Changed;
260}
261
262} // namespace sandboxir
263
264} // namespace llvm
265