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