1//===- Legality.cpp -------------------------------------------------------===//
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/Legality.h"
10#include "llvm/SandboxIR/Instruction.h"
11#include "llvm/SandboxIR/Operator.h"
12#include "llvm/SandboxIR/Utils.h"
13#include "llvm/SandboxIR/Value.h"
14#include "llvm/Support/Debug.h"
15#include "llvm/Transforms/Vectorize/SandboxVectorizer/InstrMaps.h"
16#include "llvm/Transforms/Vectorize/SandboxVectorizer/VecUtils.h"
17
18namespace llvm::sandboxir {
19
20#ifndef NDEBUG
21void ShuffleMask::dump() const {
22 print(dbgs());
23 dbgs() << "\n";
24}
25
26void LegalityResult::dump() const {
27 print(dbgs());
28 dbgs() << "\n";
29}
30#endif // NDEBUG
31
32std::optional<ResultReason>
33LegalityAnalysis::notVectorizableBasedOnOpcodesAndTypes(
34 ArrayRef<Value *> Bndl) {
35 auto *I0 = cast<Instruction>(Val: Bndl[0]);
36 auto Opcode = I0->getOpcode();
37 // If they have different opcodes, then we cannot form a vector (for now).
38 if (any_of(Range: drop_begin(RangeOrContainer&: Bndl), P: [Opcode](Value *V) {
39 return cast<Instruction>(Val: V)->getOpcode() != Opcode;
40 }))
41 return ResultReason::DiffOpcodes;
42
43 // If not the same scalar type, Pack. This will accept scalars and vectors as
44 // long as the element type is the same.
45 Type *ElmTy0 = VecUtils::getElementType(Ty: Utils::getExpectedType(V: I0));
46 if (any_of(Range: drop_begin(RangeOrContainer&: Bndl), P: [ElmTy0](Value *V) {
47 return VecUtils::getElementType(Ty: Utils::getExpectedType(V)) != ElmTy0;
48 }))
49 return ResultReason::DiffTypes;
50
51 // TODO: Allow vectorization of instrs with different flags as long as we
52 // change them to the least common one.
53 // For now pack if differnt FastMathFlags.
54 if (isa<FPMathOperator>(Val: I0)) {
55 FastMathFlags FMF0 = cast<Instruction>(Val: Bndl[0])->getFastMathFlags();
56 if (any_of(Range: drop_begin(RangeOrContainer&: Bndl), P: [FMF0](auto *V) {
57 return cast<Instruction>(V)->getFastMathFlags() != FMF0;
58 }))
59 return ResultReason::DiffMathFlags;
60 }
61
62 // TODO: Allow vectorization by using common flags.
63 // For now Pack if they don't have the same wrap flags.
64 bool CanHaveWrapFlags =
65 isa<OverflowingBinaryOperator>(Val: I0) || isa<TruncInst>(Val: I0);
66 if (CanHaveWrapFlags) {
67 bool NUW0 = I0->hasNoUnsignedWrap();
68 bool NSW0 = I0->hasNoSignedWrap();
69 if (any_of(Range: drop_begin(RangeOrContainer&: Bndl), P: [NUW0, NSW0](auto *V) {
70 return cast<Instruction>(V)->hasNoUnsignedWrap() != NUW0 ||
71 cast<Instruction>(V)->hasNoSignedWrap() != NSW0;
72 })) {
73 return ResultReason::DiffWrapFlags;
74 }
75 }
76
77 // Now we need to do further checks for specific opcodes.
78 switch (Opcode) {
79 case Instruction::Opcode::ZExt:
80 case Instruction::Opcode::SExt:
81 case Instruction::Opcode::FPToUI:
82 case Instruction::Opcode::FPToSI:
83 case Instruction::Opcode::FPExt:
84 case Instruction::Opcode::PtrToAddr:
85 case Instruction::Opcode::PtrToInt:
86 case Instruction::Opcode::IntToPtr:
87 case Instruction::Opcode::SIToFP:
88 case Instruction::Opcode::UIToFP:
89 case Instruction::Opcode::Trunc:
90 case Instruction::Opcode::FPTrunc:
91 case Instruction::Opcode::BitCast: {
92 // We have already checked that they are of the same opcode.
93 assert(all_of(Bndl,
94 [Opcode](Value *V) {
95 return cast<Instruction>(V)->getOpcode() == Opcode;
96 }) &&
97 "Different opcodes, should have early returned!");
98 // But for these opcodes we should also check the operand type.
99 Type *FromTy0 = Utils::getExpectedType(V: I0->getOperand(OpIdx: 0));
100 if (any_of(Range: drop_begin(RangeOrContainer&: Bndl), P: [FromTy0](Value *V) {
101 return Utils::getExpectedType(V: cast<User>(Val: V)->getOperand(OpIdx: 0)) !=
102 FromTy0;
103 }))
104 return ResultReason::DiffTypes;
105 return std::nullopt;
106 }
107 case Instruction::Opcode::FCmp:
108 case Instruction::Opcode::ICmp: {
109 // We need the same predicate..
110 auto Pred0 = cast<CmpInst>(Val: I0)->getPredicate();
111 bool Same = all_of(Range&: Bndl, P: [Pred0](Value *V) {
112 return cast<CmpInst>(Val: V)->getPredicate() == Pred0;
113 });
114 if (Same)
115 return std::nullopt;
116 return ResultReason::DiffOpcodes;
117 }
118 case Instruction::Opcode::Select: {
119 auto *Sel0 = cast<SelectInst>(Val: Bndl[0]);
120 auto *Cond0 = Sel0->getCondition();
121 if (VecUtils::getNumLanes(V: Cond0) != VecUtils::getNumLanes(V: Sel0))
122 // TODO: For now we don't vectorize if the lanes in the condition don't
123 // match those of the select instruction.
124 return ResultReason::Unimplemented;
125 return std::nullopt;
126 }
127 case Instruction::Opcode::FNeg:
128 case Instruction::Opcode::Add:
129 case Instruction::Opcode::FAdd:
130 case Instruction::Opcode::Sub:
131 case Instruction::Opcode::FSub:
132 case Instruction::Opcode::Mul:
133 case Instruction::Opcode::FMul:
134 case Instruction::Opcode::FRem:
135 case Instruction::Opcode::UDiv:
136 case Instruction::Opcode::SDiv:
137 case Instruction::Opcode::FDiv:
138 case Instruction::Opcode::URem:
139 case Instruction::Opcode::SRem:
140 case Instruction::Opcode::Shl:
141 case Instruction::Opcode::LShr:
142 case Instruction::Opcode::AShr:
143 case Instruction::Opcode::And:
144 case Instruction::Opcode::Or:
145 case Instruction::Opcode::Xor:
146 return std::nullopt;
147 case Instruction::Opcode::Load:
148 if (VecUtils::areConsecutive<LoadInst>(Bndl, SE, DL))
149 return std::nullopt;
150 return ResultReason::NotConsecutive;
151 case Instruction::Opcode::Store:
152 if (VecUtils::areConsecutive<StoreInst>(Bndl, SE, DL))
153 return std::nullopt;
154 return ResultReason::NotConsecutive;
155 case Instruction::Opcode::PHI:
156 return ResultReason::Unimplemented;
157 case Instruction::Opcode::Opaque:
158 return ResultReason::Unimplemented;
159 case Instruction::Opcode::Br:
160 case Instruction::Opcode::Ret:
161 case Instruction::Opcode::AddrSpaceCast:
162 case Instruction::Opcode::InsertElement:
163 case Instruction::Opcode::InsertValue:
164 case Instruction::Opcode::ExtractElement:
165 case Instruction::Opcode::ExtractValue:
166 case Instruction::Opcode::ShuffleVector:
167 case Instruction::Opcode::Call:
168 case Instruction::Opcode::GetElementPtr:
169 case Instruction::Opcode::Switch:
170 return ResultReason::Unimplemented;
171 case Instruction::Opcode::VAArg:
172 case Instruction::Opcode::Freeze:
173 case Instruction::Opcode::Fence:
174 case Instruction::Opcode::Invoke:
175 case Instruction::Opcode::CallBr:
176 case Instruction::Opcode::LandingPad:
177 case Instruction::Opcode::CatchPad:
178 case Instruction::Opcode::CleanupPad:
179 case Instruction::Opcode::CatchRet:
180 case Instruction::Opcode::CleanupRet:
181 case Instruction::Opcode::Resume:
182 case Instruction::Opcode::CatchSwitch:
183 case Instruction::Opcode::AtomicRMW:
184 case Instruction::Opcode::AtomicCmpXchg:
185 case Instruction::Opcode::Alloca:
186 case Instruction::Opcode::Unreachable:
187 return ResultReason::Infeasible;
188 }
189
190 return std::nullopt;
191}
192
193CollectDescr
194LegalityAnalysis::getHowToCollectValues(ArrayRef<Value *> Bndl) const {
195 SmallVector<CollectDescr::ExtractElementDescr, 4> Vec;
196 Vec.reserve(N: Bndl.size());
197 for (auto [Elm, V] : enumerate(First&: Bndl)) {
198 if (auto *VecOp = IMaps.getVectorForOrig(Orig: V)) {
199 // If there is a vector containing `V`, then get the lane it came from.
200 std::optional<int> ExtractIdxOpt = IMaps.getOrigLane(Vec: VecOp, Orig: V);
201 // This could be a vector, like <2 x float> in which case the mask needs
202 // to enumerate all lanes.
203 for (unsigned Ln = 0, Lanes = VecUtils::getNumLanes(V); Ln != Lanes; ++Ln)
204 Vec.emplace_back(Args&: VecOp, Args: ExtractIdxOpt ? *ExtractIdxOpt + Ln : -1);
205 } else {
206 Vec.emplace_back(Args: V);
207 }
208 }
209 return CollectDescr(std::move(Vec));
210}
211
212const LegalityResult &LegalityAnalysis::canVectorize(ArrayRef<Value *> Bndl,
213 bool SkipScheduling) {
214 // If Bndl contains values other than instructions, we need to Pack.
215 if (any_of(Range&: Bndl, P: [](auto *V) { return !isa<Instruction>(V); }))
216 return createLegalityResult<Pack>(Args: ResultReason::NotInstructions);
217 // Pack if not in the same BB.
218 auto *BB = cast<Instruction>(Val: Bndl[0])->getParent();
219 if (any_of(Range: drop_begin(RangeOrContainer&: Bndl),
220 P: [BB](auto *V) { return cast<Instruction>(V)->getParent() != BB; }))
221 return createLegalityResult<Pack>(Args: ResultReason::DiffBBs);
222 // Pack if instructions repeat, i.e., require some sort of broadcast.
223 SmallPtrSet<Value *, 8> Unique(llvm::from_range, Bndl);
224 if (Unique.size() != Bndl.size())
225 return createLegalityResult<Pack>(Args: ResultReason::RepeatedInstrs);
226
227 auto CollectDescrs = getHowToCollectValues(Bndl);
228 if (CollectDescrs.hasVectorInputs()) {
229 if (auto ValueShuffleOpt = CollectDescrs.getSingleInput()) {
230 auto [Vec, Mask] = *ValueShuffleOpt;
231 if (Mask.isIdentity())
232 return createLegalityResult<DiamondReuse>(Args&: Vec);
233 return createLegalityResult<DiamondReuseWithShuffle>(Args&: Vec, Args&: Mask);
234 }
235 return createLegalityResult<DiamondReuseMultiInput>(
236 Args: std::move(CollectDescrs));
237 }
238
239 if (auto ReasonOpt = notVectorizableBasedOnOpcodesAndTypes(Bndl))
240 return createLegalityResult<Pack>(Args&: *ReasonOpt);
241
242 if (!SkipScheduling) {
243 // TODO: Try to remove the IBndl vector.
244 SmallVector<Instruction *, 8> IBndl;
245 IBndl.reserve(N: Bndl.size());
246 for (auto *V : Bndl)
247 IBndl.push_back(Elt: cast<Instruction>(Val: V));
248 if (!Sched.trySchedule(Instrs: IBndl))
249 return createLegalityResult<Pack>(Args: ResultReason::CantSchedule);
250 }
251
252 return createLegalityResult<Widen>();
253}
254
255void LegalityAnalysis::clear() {
256 Sched.clear();
257 IMaps.clear();
258}
259} // namespace llvm::sandboxir
260