1//===- ScalarizeMaskedMemIntrin.cpp - Scalarize unsupported masked mem ----===//
2// intrinsics
3//
4// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
5// See https://llvm.org/LICENSE.txt for license information.
6// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
7//
8//===----------------------------------------------------------------------===//
9//
10// This pass replaces masked memory intrinsics - when unsupported by the target
11// - with a chain of basic blocks, that deal with the elements one-by-one if the
12// appropriate mask bit is set.
13//
14//===----------------------------------------------------------------------===//
15
16#include "llvm/Transforms/Scalar/ScalarizeMaskedMemIntrin.h"
17#include "llvm/ADT/Twine.h"
18#include "llvm/Analysis/DomTreeUpdater.h"
19#include "llvm/Analysis/TargetTransformInfo.h"
20#include "llvm/Analysis/VectorUtils.h"
21#include "llvm/IR/BasicBlock.h"
22#include "llvm/IR/Constant.h"
23#include "llvm/IR/Constants.h"
24#include "llvm/IR/DerivedTypes.h"
25#include "llvm/IR/Dominators.h"
26#include "llvm/IR/Function.h"
27#include "llvm/IR/IRBuilder.h"
28#include "llvm/IR/Instruction.h"
29#include "llvm/IR/Instructions.h"
30#include "llvm/IR/IntrinsicInst.h"
31#include "llvm/IR/Metadata.h"
32#include "llvm/IR/ProfDataUtils.h"
33#include "llvm/IR/Type.h"
34#include "llvm/IR/Value.h"
35#include "llvm/InitializePasses.h"
36#include "llvm/Pass.h"
37#include "llvm/Support/Casting.h"
38#include "llvm/Transforms/Scalar.h"
39#include "llvm/Transforms/Utils/BasicBlockUtils.h"
40#include <cassert>
41#include <optional>
42
43using namespace llvm;
44
45#define DEBUG_TYPE "scalarize-masked-mem-intrin"
46
47namespace {
48
49class ScalarizeMaskedMemIntrinLegacyPass : public FunctionPass {
50public:
51 static char ID; // Pass identification, replacement for typeid
52
53 explicit ScalarizeMaskedMemIntrinLegacyPass() : FunctionPass(ID) {
54 initializeScalarizeMaskedMemIntrinLegacyPassPass(
55 *PassRegistry::getPassRegistry());
56 }
57
58 bool runOnFunction(Function &F) override;
59
60 StringRef getPassName() const override {
61 return "Scalarize Masked Memory Intrinsics";
62 }
63
64 void getAnalysisUsage(AnalysisUsage &AU) const override {
65 AU.addRequired<TargetTransformInfoWrapperPass>();
66 AU.addPreserved<DominatorTreeWrapperPass>();
67 }
68};
69
70} // end anonymous namespace
71
72static bool optimizeBlock(BasicBlock &BB, bool &ModifiedDT,
73 const TargetTransformInfo &TTI, const DataLayout &DL,
74 bool HasBranchDivergence, DomTreeUpdater *DTU);
75static bool optimizeCallInst(CallInst *CI, bool &ModifiedDT,
76 const TargetTransformInfo &TTI,
77 const DataLayout &DL, bool HasBranchDivergence,
78 DomTreeUpdater *DTU);
79
80char ScalarizeMaskedMemIntrinLegacyPass::ID = 0;
81
82INITIALIZE_PASS_BEGIN(ScalarizeMaskedMemIntrinLegacyPass, DEBUG_TYPE,
83 "Scalarize unsupported masked memory intrinsics", false,
84 false)
85INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
86INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
87INITIALIZE_PASS_END(ScalarizeMaskedMemIntrinLegacyPass, DEBUG_TYPE,
88 "Scalarize unsupported masked memory intrinsics", false,
89 false)
90
91FunctionPass *llvm::createScalarizeMaskedMemIntrinLegacyPass() {
92 return new ScalarizeMaskedMemIntrinLegacyPass();
93}
94
95static bool isConstantIntVector(Value *Mask) {
96 Constant *C = dyn_cast<Constant>(Val: Mask);
97 if (!C)
98 return false;
99
100 unsigned NumElts = cast<FixedVectorType>(Val: Mask->getType())->getNumElements();
101 for (unsigned i = 0; i != NumElts; ++i) {
102 Constant *CElt = C->getAggregateElement(Elt: i);
103 if (!CElt || !isa<ConstantInt>(Val: CElt))
104 return false;
105 }
106
107 return true;
108}
109
110static unsigned adjustForEndian(const DataLayout &DL, unsigned VectorWidth,
111 unsigned Idx) {
112 return DL.isBigEndian() ? VectorWidth - 1 - Idx : Idx;
113}
114
115static void copyMemCacheHint(Instruction &Dest, const Instruction &Source,
116 unsigned SourcePtrOperand,
117 unsigned DestPtrOperand) {
118 MDNode *CacheHint = Source.getMetadata(KindID: LLVMContext::MD_mem_cache_hint);
119 // These intrinsics have a single memory operand.
120 if (!CacheHint || CacheHint->getNumOperands() != 2)
121 return;
122
123 auto *OperandNo = mdconst::extract<ConstantInt>(MD: CacheHint->getOperand(I: 0));
124 if (OperandNo->getZExtValue() != SourcePtrOperand)
125 return;
126
127 Metadata *DestOperandNo = ConstantAsMetadata::get(
128 C: ConstantInt::get(Ty: Type::getInt32Ty(C&: Dest.getContext()), V: DestPtrOperand));
129 Dest.setMetadata(KindID: LLVMContext::MD_mem_cache_hint,
130 Node: MDNode::get(Context&: Dest.getContext(),
131 MDs: {DestOperandNo, CacheHint->getOperand(I: 1)}));
132}
133
134static void copyMetadataForMemoryAccess(
135 Instruction &Dest, const Instruction &Source, const DataLayout &DL,
136 unsigned SourcePtrOperand, unsigned DestPtrOperand, Type *AccessType,
137 bool IsWholeAccess, std::optional<size_t> ByteOffset) {
138 // Only propagate metadata that is valid on each constituent memory access.
139 // In particular, do not copy metadata whose meaning is tied to the call,
140 // such as !prof or !callsite.
141 Dest.copyMetadata(SrcInst: Source,
142 WL: {LLVMContext::MD_nontemporal,
143 LLVMContext::MD_mem_parallel_loop_access,
144 LLVMContext::MD_access_group, LLVMContext::MD_annotation,
145 LLVMContext::MD_nosanitize, LLVMContext::MD_mmra});
146
147 AAMDNodes AANodes = Source.getAAMetadata();
148 if (IsWholeAccess)
149 Dest.setAAMetadata(AANodes);
150 else if (ByteOffset)
151 Dest.setAAMetadata(AANodes.adjustForAccess(Offset: *ByteOffset, AccessTy: AccessType, DL));
152 else {
153 // The packed address is runtime-dependent. The other AA metadata remains
154 // applicable, but !tbaa.struct cannot be adjusted to a known byte range.
155 AANodes.TBAAStruct = nullptr;
156 Dest.setAAMetadata(AANodes);
157 }
158 copyMemCacheHint(Dest, Source, SourcePtrOperand, DestPtrOperand);
159}
160
161static void copyMetadataForScalarizedLoad(LoadInst &Dest,
162 const Instruction &Source,
163 const DataLayout &DL,
164 unsigned SourcePtrOperand,
165 std::optional<size_t> ByteOffset) {
166 copyMetadataForMemoryAccess(Dest, Source, DL, SourcePtrOperand,
167 DestPtrOperand: Dest.getPointerOperandIndex(), AccessType: Dest.getType(),
168 IsWholeAccess: Dest.getType() == Source.getType(), ByteOffset);
169
170 // !range applies element-wise to vectors, so the same range describes each
171 // scalar result. The other metadata here also describes the loaded result.
172 Dest.copyMetadata(SrcInst: Source, WL: {LLVMContext::MD_fpmath, LLVMContext::MD_range,
173 LLVMContext::MD_invariant_load});
174}
175
176static void copyMetadataForScalarizedStore(StoreInst &Dest,
177 const Instruction &Source,
178 const DataLayout &DL,
179 unsigned SourcePtrOperand,
180 std::optional<size_t> ByteOffset) {
181 copyMetadataForMemoryAccess(
182 Dest, Source, DL, SourcePtrOperand, DestPtrOperand: Dest.getPointerOperandIndex(),
183 AccessType: Dest.getValueOperand()->getType(),
184 IsWholeAccess: Dest.getValueOperand()->getType() == Source.getOperand(i: 0)->getType(),
185 ByteOffset);
186}
187
188// Translate a masked load intrinsic like
189// <16 x i32 > @llvm.masked.load( <16 x i32>* %addr,
190// <16 x i1> %mask, <16 x i32> %passthru)
191// to a chain of basic blocks, with loading element one-by-one if
192// the appropriate mask bit is set
193//
194// %1 = bitcast i8* %addr to i32*
195// %2 = extractelement <16 x i1> %mask, i32 0
196// br i1 %2, label %cond.load, label %else
197//
198// cond.load: ; preds = %0
199// %3 = getelementptr i32* %1, i32 0
200// %4 = load i32* %3
201// %5 = insertelement <16 x i32> %passthru, i32 %4, i32 0
202// br label %else
203//
204// else: ; preds = %0, %cond.load
205// %res.phi.else = phi <16 x i32> [ %5, %cond.load ], [ poison, %0 ]
206// %6 = extractelement <16 x i1> %mask, i32 1
207// br i1 %6, label %cond.load1, label %else2
208//
209// cond.load1: ; preds = %else
210// %7 = getelementptr i32* %1, i32 1
211// %8 = load i32* %7
212// %9 = insertelement <16 x i32> %res.phi.else, i32 %8, i32 1
213// br label %else2
214//
215// else2: ; preds = %else, %cond.load1
216// %res.phi.else3 = phi <16 x i32> [ %9, %cond.load1 ], [ %res.phi.else, %else
217// ] %10 = extractelement <16 x i1> %mask, i32 2 br i1 %10, label %cond.load4,
218// label %else5
219//
220static void scalarizeMaskedLoad(const DataLayout &DL, bool HasBranchDivergence,
221 CallInst *CI, DomTreeUpdater *DTU,
222 bool &ModifiedDT) {
223 Value *Ptr = CI->getArgOperand(i: 0);
224 Value *Mask = CI->getArgOperand(i: 1);
225 Value *Src0 = CI->getArgOperand(i: 2);
226
227 const Align AlignVal = CI->getParamAlign(ArgNo: 0).valueOrOne();
228 VectorType *VecType = cast<FixedVectorType>(Val: CI->getType());
229
230 Type *EltTy = VecType->getElementType();
231
232 IRBuilder<> Builder(CI->getContext());
233 Instruction *InsertPt = CI;
234 BasicBlock *IfBlock = CI->getParent();
235
236 Builder.SetInsertPoint(InsertPt);
237 Builder.SetCurrentDebugLocation(CI->getDebugLoc());
238
239 // Short-cut if the mask is all-true.
240 if (isa<Constant>(Val: Mask) && cast<Constant>(Val: Mask)->isAllOnesValue()) {
241 LoadInst *NewI = Builder.CreateAlignedLoad(Ty: VecType, Ptr, Align: AlignVal);
242 copyMetadataForScalarizedLoad(Dest&: *NewI, Source: *CI, DL, /*SourcePtrOperand=*/0,
243 ByteOffset: std::nullopt);
244 NewI->takeName(V: CI);
245 CI->replaceAllUsesWith(V: NewI);
246 CI->eraseFromParent();
247 return;
248 }
249
250 // Adjust alignment for the scalar instruction.
251 const Align AdjustedAlignVal =
252 commonAlignment(A: AlignVal, Offset: EltTy->getPrimitiveSizeInBits() / 8);
253 unsigned VectorWidth = cast<FixedVectorType>(Val: VecType)->getNumElements();
254
255 // The result vector
256 Value *VResult = Src0;
257
258 if (isConstantIntVector(Mask)) {
259 for (unsigned Idx = 0; Idx < VectorWidth; ++Idx) {
260 if (cast<Constant>(Val: Mask)->getAggregateElement(Elt: Idx)->isNullValue())
261 continue;
262 Value *Gep = Builder.CreateConstInBoundsGEP1_32(Ty: EltTy, Ptr, Idx0: Idx);
263 LoadInst *Load = Builder.CreateAlignedLoad(Ty: EltTy, Ptr: Gep, Align: AdjustedAlignVal);
264 copyMetadataForScalarizedLoad(
265 Dest&: *Load, Source: *CI, DL, /*SourcePtrOperand=*/0,
266 ByteOffset: Idx * DL.getTypeAllocSize(Ty: EltTy).getFixedValue());
267 VResult = Builder.CreateInsertElement(Vec: VResult, NewElt: Load, Idx);
268 }
269 CI->replaceAllUsesWith(V: VResult);
270 CI->eraseFromParent();
271 return;
272 }
273
274 // Optimize the case where the "masked load" is a predicated load - that is,
275 // where the mask is the splat of a non-constant scalar boolean. In that case,
276 // use that splated value as the guard on a conditional vector load.
277 if (isSplatValue(V: Mask, /*Index=*/0)) {
278 Value *Predicate = Builder.CreateExtractElement(Vec: Mask, Idx: uint64_t(0ull),
279 Name: Mask->getName() + ".first");
280 // We mark the branch weights as explicitly unknown given they would only
281 // be derivable from the mask which we do not have VP information for.
282 Instruction *ThenTerm =
283 SplitBlockAndInsertIfThen(Cond: Predicate, SplitBefore: InsertPt, /*Unreachable=*/false,
284 BranchWeights: getExplicitlyUnknownBranchWeightsIfProfiled(
285 F&: *CI->getFunction(), DEBUG_TYPE),
286 DTU);
287
288 BasicBlock *CondBlock = ThenTerm->getParent();
289 CondBlock->setName("cond.load");
290 Builder.SetInsertPoint(CondBlock->getTerminator());
291 LoadInst *Load = Builder.CreateAlignedLoad(Ty: VecType, Ptr, Align: AlignVal,
292 Name: CI->getName() + ".cond.load");
293 copyMetadataForScalarizedLoad(Dest&: *Load, Source: *CI, DL, /*SourcePtrOperand=*/0,
294 ByteOffset: std::nullopt);
295
296 BasicBlock *PostLoad = ThenTerm->getSuccessor(Idx: 0);
297 Builder.SetInsertPoint(TheBB: PostLoad, IP: PostLoad->begin());
298 PHINode *Phi = Builder.CreatePHI(Ty: VecType, /*NumReservedValues=*/2);
299 Phi->addIncoming(V: Load, BB: CondBlock);
300 Phi->addIncoming(V: Src0, BB: IfBlock);
301 Phi->takeName(V: CI);
302
303 CI->replaceAllUsesWith(V: Phi);
304 CI->eraseFromParent();
305 ModifiedDT = true;
306 return;
307 }
308 // If the mask is not v1i1, use scalar bit test operations. This generates
309 // better results on X86 at least. However, don't do this on GPUs and other
310 // machines with divergence, as there each i1 needs a vector register.
311 Value *SclrMask = nullptr;
312 if (VectorWidth != 1 && !HasBranchDivergence) {
313 Type *SclrMaskTy = Builder.getIntNTy(N: VectorWidth);
314 SclrMask = Builder.CreateBitCast(V: Mask, DestTy: SclrMaskTy, Name: "scalar_mask");
315 }
316
317 for (unsigned Idx = 0; Idx < VectorWidth; ++Idx) {
318 // Fill the "else" block, created in the previous iteration
319 //
320 // %res.phi.else3 = phi <16 x i32> [ %11, %cond.load1 ], [ %res.phi.else,
321 // %else ] %mask_1 = and i16 %scalar_mask, i32 1 << Idx %cond = icmp ne i16
322 // %mask_1, 0 br i1 %mask_1, label %cond.load, label %else
323 //
324 // On GPUs, use
325 // %cond = extrectelement %mask, Idx
326 // instead
327 Value *Predicate;
328 if (SclrMask != nullptr) {
329 Value *Mask = Builder.getInt(AI: APInt::getOneBitSet(
330 numBits: VectorWidth, BitNo: adjustForEndian(DL, VectorWidth, Idx)));
331 Predicate = Builder.CreateICmpNE(LHS: Builder.CreateAnd(LHS: SclrMask, RHS: Mask),
332 RHS: Builder.getIntN(N: VectorWidth, C: 0));
333 } else {
334 Predicate = Builder.CreateExtractElement(Vec: Mask, Idx);
335 }
336
337 // Create "cond" block
338 //
339 // %EltAddr = getelementptr i32* %1, i32 0
340 // %Elt = load i32* %EltAddr
341 // VResult = insertelement <16 x i32> VResult, i32 %Elt, i32 Idx
342 //
343 // We mark the branch weights as explicitly unknown given they would only
344 // be derivable from the mask which we do not have VP information for.
345 Instruction *ThenTerm =
346 SplitBlockAndInsertIfThen(Cond: Predicate, SplitBefore: InsertPt, /*Unreachable=*/false,
347 BranchWeights: getExplicitlyUnknownBranchWeightsIfProfiled(
348 F&: *CI->getFunction(), DEBUG_TYPE),
349 DTU);
350
351 BasicBlock *CondBlock = ThenTerm->getParent();
352 CondBlock->setName("cond.load");
353
354 Builder.SetInsertPoint(CondBlock->getTerminator());
355 Value *Gep = Builder.CreateConstInBoundsGEP1_32(Ty: EltTy, Ptr, Idx0: Idx);
356 LoadInst *Load = Builder.CreateAlignedLoad(Ty: EltTy, Ptr: Gep, Align: AdjustedAlignVal);
357 copyMetadataForScalarizedLoad(
358 Dest&: *Load, Source: *CI, DL, /*SourcePtrOperand=*/0,
359 ByteOffset: Idx * DL.getTypeAllocSize(Ty: EltTy).getFixedValue());
360 Value *NewVResult = Builder.CreateInsertElement(Vec: VResult, NewElt: Load, Idx);
361
362 // Create "else" block, fill it in the next iteration
363 BasicBlock *NewIfBlock = ThenTerm->getSuccessor(Idx: 0);
364 NewIfBlock->setName("else");
365 BasicBlock *PrevIfBlock = IfBlock;
366 IfBlock = NewIfBlock;
367
368 // Create the phi to join the new and previous value.
369 Builder.SetInsertPoint(TheBB: NewIfBlock, IP: NewIfBlock->begin());
370 PHINode *Phi = Builder.CreatePHI(Ty: VecType, NumReservedValues: 2, Name: "res.phi.else");
371 Phi->addIncoming(V: NewVResult, BB: CondBlock);
372 Phi->addIncoming(V: VResult, BB: PrevIfBlock);
373 VResult = Phi;
374 }
375
376 CI->replaceAllUsesWith(V: VResult);
377 CI->eraseFromParent();
378
379 ModifiedDT = true;
380}
381
382// Translate a masked store intrinsic, like
383// void @llvm.masked.store(<16 x i32> %src, <16 x i32>* %addr,
384// <16 x i1> %mask)
385// to a chain of basic blocks, that stores element one-by-one if
386// the appropriate mask bit is set
387//
388// %1 = bitcast i8* %addr to i32*
389// %2 = extractelement <16 x i1> %mask, i32 0
390// br i1 %2, label %cond.store, label %else
391//
392// cond.store: ; preds = %0
393// %3 = extractelement <16 x i32> %val, i32 0
394// %4 = getelementptr i32* %1, i32 0
395// store i32 %3, i32* %4
396// br label %else
397//
398// else: ; preds = %0, %cond.store
399// %5 = extractelement <16 x i1> %mask, i32 1
400// br i1 %5, label %cond.store1, label %else2
401//
402// cond.store1: ; preds = %else
403// %6 = extractelement <16 x i32> %val, i32 1
404// %7 = getelementptr i32* %1, i32 1
405// store i32 %6, i32* %7
406// br label %else2
407// . . .
408static void scalarizeMaskedStore(const DataLayout &DL, bool HasBranchDivergence,
409 CallInst *CI, DomTreeUpdater *DTU,
410 bool &ModifiedDT) {
411 Value *Src = CI->getArgOperand(i: 0);
412 Value *Ptr = CI->getArgOperand(i: 1);
413 Value *Mask = CI->getArgOperand(i: 2);
414
415 const Align AlignVal = CI->getParamAlign(ArgNo: 1).valueOrOne();
416 auto *VecType = cast<VectorType>(Val: Src->getType());
417
418 Type *EltTy = VecType->getElementType();
419
420 IRBuilder<> Builder(CI->getContext());
421 Instruction *InsertPt = CI;
422 Builder.SetInsertPoint(InsertPt);
423 Builder.SetCurrentDebugLocation(CI->getDebugLoc());
424
425 // Short-cut if the mask is all-true.
426 if (isa<Constant>(Val: Mask) && cast<Constant>(Val: Mask)->isAllOnesValue()) {
427 StoreInst *Store = Builder.CreateAlignedStore(Val: Src, Ptr, Align: AlignVal);
428 Store->takeName(V: CI);
429 copyMetadataForScalarizedStore(Dest&: *Store, Source: *CI, DL, /*SourcePtrOperand=*/1,
430 ByteOffset: std::nullopt);
431 // This is a one-to-one replacement, so the assignment link remains valid.
432 Store->copyMetadata(SrcInst: *CI, WL: LLVMContext::MD_DIAssignID);
433 CI->eraseFromParent();
434 return;
435 }
436
437 // Adjust alignment for the scalar instruction.
438 const Align AdjustedAlignVal =
439 commonAlignment(A: AlignVal, Offset: EltTy->getPrimitiveSizeInBits() / 8);
440 unsigned VectorWidth = cast<FixedVectorType>(Val: VecType)->getNumElements();
441
442 if (isConstantIntVector(Mask)) {
443 for (unsigned Idx = 0; Idx < VectorWidth; ++Idx) {
444 if (cast<Constant>(Val: Mask)->getAggregateElement(Elt: Idx)->isNullValue())
445 continue;
446 Value *OneElt = Builder.CreateExtractElement(Vec: Src, Idx);
447 Value *Gep = Builder.CreateConstInBoundsGEP1_32(Ty: EltTy, Ptr, Idx0: Idx);
448 StoreInst *Store =
449 Builder.CreateAlignedStore(Val: OneElt, Ptr: Gep, Align: AdjustedAlignVal);
450 copyMetadataForScalarizedStore(
451 Dest&: *Store, Source: *CI, DL, /*SourcePtrOperand=*/1,
452 ByteOffset: Idx * DL.getTypeAllocSize(Ty: EltTy).getFixedValue());
453 }
454 CI->eraseFromParent();
455 return;
456 }
457
458 // Optimize the case where the "masked store" is a predicated store - that is,
459 // when the mask is the splat of a non-constant scalar boolean. In that case,
460 // optimize to a conditional store.
461 if (isSplatValue(V: Mask, /*Index=*/0)) {
462 Value *Predicate = Builder.CreateExtractElement(Vec: Mask, Idx: uint64_t(0ull),
463 Name: Mask->getName() + ".first");
464 // We mark the branch weights as explicitly unknown given they would only
465 // be derivable from the mask which we do not have VP information for.
466 Instruction *ThenTerm =
467 SplitBlockAndInsertIfThen(Cond: Predicate, SplitBefore: InsertPt, /*Unreachable=*/false,
468 BranchWeights: getExplicitlyUnknownBranchWeightsIfProfiled(
469 F&: *CI->getFunction(), DEBUG_TYPE),
470 DTU);
471 BasicBlock *CondBlock = ThenTerm->getParent();
472 CondBlock->setName("cond.store");
473 Builder.SetInsertPoint(CondBlock->getTerminator());
474
475 StoreInst *Store = Builder.CreateAlignedStore(Val: Src, Ptr, Align: AlignVal);
476 Store->takeName(V: CI);
477 copyMetadataForScalarizedStore(Dest&: *Store, Source: *CI, DL, /*SourcePtrOperand=*/1,
478 ByteOffset: std::nullopt);
479 // This is a one-to-one replacement, so the assignment link remains valid.
480 Store->copyMetadata(SrcInst: *CI, WL: LLVMContext::MD_DIAssignID);
481
482 CI->eraseFromParent();
483 ModifiedDT = true;
484 return;
485 }
486
487 // If the mask is not v1i1, use scalar bit test operations. This generates
488 // better results on X86 at least. However, don't do this on GPUs or other
489 // machines with branch divergence, as there each i1 takes up a register.
490 Value *SclrMask = nullptr;
491 if (VectorWidth != 1 && !HasBranchDivergence) {
492 Type *SclrMaskTy = Builder.getIntNTy(N: VectorWidth);
493 SclrMask = Builder.CreateBitCast(V: Mask, DestTy: SclrMaskTy, Name: "scalar_mask");
494 }
495
496 for (unsigned Idx = 0; Idx < VectorWidth; ++Idx) {
497 // Fill the "else" block, created in the previous iteration
498 //
499 // %mask_1 = and i16 %scalar_mask, i32 1 << Idx
500 // %cond = icmp ne i16 %mask_1, 0
501 // br i1 %mask_1, label %cond.store, label %else
502 //
503 // On GPUs, use
504 // %cond = extrectelement %mask, Idx
505 // instead
506 Value *Predicate;
507 if (SclrMask != nullptr) {
508 Value *Mask = Builder.getInt(AI: APInt::getOneBitSet(
509 numBits: VectorWidth, BitNo: adjustForEndian(DL, VectorWidth, Idx)));
510 Predicate = Builder.CreateICmpNE(LHS: Builder.CreateAnd(LHS: SclrMask, RHS: Mask),
511 RHS: Builder.getIntN(N: VectorWidth, C: 0));
512 } else {
513 Predicate = Builder.CreateExtractElement(Vec: Mask, Idx);
514 }
515
516 // Create "cond" block
517 //
518 // %OneElt = extractelement <16 x i32> %Src, i32 Idx
519 // %EltAddr = getelementptr i32* %1, i32 0
520 // %store i32 %OneElt, i32* %EltAddr
521 //
522 // We mark the branch weights as explicitly unknown given they would only
523 // be derivable from the mask which we do not have VP information for.
524 Instruction *ThenTerm =
525 SplitBlockAndInsertIfThen(Cond: Predicate, SplitBefore: InsertPt, /*Unreachable=*/false,
526 BranchWeights: getExplicitlyUnknownBranchWeightsIfProfiled(
527 F&: *CI->getFunction(), DEBUG_TYPE),
528 DTU);
529
530 BasicBlock *CondBlock = ThenTerm->getParent();
531 CondBlock->setName("cond.store");
532
533 Builder.SetInsertPoint(CondBlock->getTerminator());
534 Value *OneElt = Builder.CreateExtractElement(Vec: Src, Idx);
535 Value *Gep = Builder.CreateConstInBoundsGEP1_32(Ty: EltTy, Ptr, Idx0: Idx);
536 StoreInst *Store =
537 Builder.CreateAlignedStore(Val: OneElt, Ptr: Gep, Align: AdjustedAlignVal);
538 copyMetadataForScalarizedStore(
539 Dest&: *Store, Source: *CI, DL, /*SourcePtrOperand=*/1,
540 ByteOffset: Idx * DL.getTypeAllocSize(Ty: EltTy).getFixedValue());
541
542 // Create "else" block, fill it in the next iteration
543 BasicBlock *NewIfBlock = ThenTerm->getSuccessor(Idx: 0);
544 NewIfBlock->setName("else");
545
546 Builder.SetInsertPoint(TheBB: NewIfBlock, IP: NewIfBlock->begin());
547 }
548 CI->eraseFromParent();
549
550 ModifiedDT = true;
551}
552
553// Translate a masked gather intrinsic like
554// <16 x i32 > @llvm.masked.gather.v16i32( <16 x i32*> %Ptrs, i32 4,
555// <16 x i1> %Mask, <16 x i32> %Src)
556// to a chain of basic blocks, with loading element one-by-one if
557// the appropriate mask bit is set
558//
559// %Ptrs = getelementptr i32, i32* %base, <16 x i64> %ind
560// %Mask0 = extractelement <16 x i1> %Mask, i32 0
561// br i1 %Mask0, label %cond.load, label %else
562//
563// cond.load:
564// %Ptr0 = extractelement <16 x i32*> %Ptrs, i32 0
565// %Load0 = load i32, i32* %Ptr0, align 4
566// %Res0 = insertelement <16 x i32> poison, i32 %Load0, i32 0
567// br label %else
568//
569// else:
570// %res.phi.else = phi <16 x i32>[%Res0, %cond.load], [poison, %0]
571// %Mask1 = extractelement <16 x i1> %Mask, i32 1
572// br i1 %Mask1, label %cond.load1, label %else2
573//
574// cond.load1:
575// %Ptr1 = extractelement <16 x i32*> %Ptrs, i32 1
576// %Load1 = load i32, i32* %Ptr1, align 4
577// %Res1 = insertelement <16 x i32> %res.phi.else, i32 %Load1, i32 1
578// br label %else2
579// . . .
580// %Result = select <16 x i1> %Mask, <16 x i32> %res.phi.select, <16 x i32> %Src
581// ret <16 x i32> %Result
582static void scalarizeMaskedGather(const DataLayout &DL,
583 bool HasBranchDivergence, CallInst *CI,
584 DomTreeUpdater *DTU, bool &ModifiedDT) {
585 Value *Ptrs = CI->getArgOperand(i: 0);
586 Value *Mask = CI->getArgOperand(i: 1);
587 Value *Src0 = CI->getArgOperand(i: 2);
588
589 auto *VecType = cast<FixedVectorType>(Val: CI->getType());
590 Type *EltTy = VecType->getElementType();
591
592 IRBuilder<> Builder(CI->getContext());
593 Instruction *InsertPt = CI;
594 BasicBlock *IfBlock = CI->getParent();
595 Builder.SetInsertPoint(InsertPt);
596 Align AlignVal = CI->getParamAlign(ArgNo: 0).valueOrOne();
597
598 Builder.SetCurrentDebugLocation(CI->getDebugLoc());
599
600 // The result vector
601 Value *VResult = Src0;
602 unsigned VectorWidth = VecType->getNumElements();
603
604 // Shorten the way if the mask is a vector of constants.
605 if (isConstantIntVector(Mask)) {
606 for (unsigned Idx = 0; Idx < VectorWidth; ++Idx) {
607 if (cast<Constant>(Val: Mask)->getAggregateElement(Elt: Idx)->isNullValue())
608 continue;
609 Value *Ptr = Builder.CreateExtractElement(Vec: Ptrs, Idx, Name: "Ptr" + Twine(Idx));
610 LoadInst *Load =
611 Builder.CreateAlignedLoad(Ty: EltTy, Ptr, Align: AlignVal, Name: "Load" + Twine(Idx));
612 copyMetadataForScalarizedLoad(Dest&: *Load, Source: *CI, DL, /*SourcePtrOperand=*/0,
613 /*ByteOffset=*/0);
614 VResult =
615 Builder.CreateInsertElement(Vec: VResult, NewElt: Load, Idx, Name: "Res" + Twine(Idx));
616 }
617 CI->replaceAllUsesWith(V: VResult);
618 CI->eraseFromParent();
619 return;
620 }
621
622 // If the mask is not v1i1, use scalar bit test operations. This generates
623 // better results on X86 at least. However, don't do this on GPUs or other
624 // machines with branch divergence, as there, each i1 takes up a register.
625 Value *SclrMask = nullptr;
626 if (VectorWidth != 1 && !HasBranchDivergence) {
627 Type *SclrMaskTy = Builder.getIntNTy(N: VectorWidth);
628 SclrMask = Builder.CreateBitCast(V: Mask, DestTy: SclrMaskTy, Name: "scalar_mask");
629 }
630
631 for (unsigned Idx = 0; Idx < VectorWidth; ++Idx) {
632 // Fill the "else" block, created in the previous iteration
633 //
634 // %Mask1 = and i16 %scalar_mask, i32 1 << Idx
635 // %cond = icmp ne i16 %mask_1, 0
636 // br i1 %Mask1, label %cond.load, label %else
637 //
638 // On GPUs, use
639 // %cond = extrectelement %mask, Idx
640 // instead
641
642 Value *Predicate;
643 if (SclrMask != nullptr) {
644 Value *Mask = Builder.getInt(AI: APInt::getOneBitSet(
645 numBits: VectorWidth, BitNo: adjustForEndian(DL, VectorWidth, Idx)));
646 Predicate = Builder.CreateICmpNE(LHS: Builder.CreateAnd(LHS: SclrMask, RHS: Mask),
647 RHS: Builder.getIntN(N: VectorWidth, C: 0));
648 } else {
649 Predicate = Builder.CreateExtractElement(Vec: Mask, Idx, Name: "Mask" + Twine(Idx));
650 }
651
652 // Create "cond" block
653 //
654 // %EltAddr = getelementptr i32* %1, i32 0
655 // %Elt = load i32* %EltAddr
656 // VResult = insertelement <16 x i32> VResult, i32 %Elt, i32 Idx
657 //
658 // We mark the branch weights as explicitly unknown given they would only
659 // be derivable from the mask which we do not have VP information for.
660 Instruction *ThenTerm =
661 SplitBlockAndInsertIfThen(Cond: Predicate, SplitBefore: InsertPt, /*Unreachable=*/false,
662 BranchWeights: getExplicitlyUnknownBranchWeightsIfProfiled(
663 F&: *CI->getFunction(), DEBUG_TYPE),
664 DTU);
665
666 BasicBlock *CondBlock = ThenTerm->getParent();
667 CondBlock->setName("cond.load");
668
669 Builder.SetInsertPoint(CondBlock->getTerminator());
670 Value *Ptr = Builder.CreateExtractElement(Vec: Ptrs, Idx, Name: "Ptr" + Twine(Idx));
671 LoadInst *Load =
672 Builder.CreateAlignedLoad(Ty: EltTy, Ptr, Align: AlignVal, Name: "Load" + Twine(Idx));
673 copyMetadataForScalarizedLoad(Dest&: *Load, Source: *CI, DL, /*SourcePtrOperand=*/0,
674 /*ByteOffset=*/0);
675 Value *NewVResult =
676 Builder.CreateInsertElement(Vec: VResult, NewElt: Load, Idx, Name: "Res" + Twine(Idx));
677
678 // Create "else" block, fill it in the next iteration
679 BasicBlock *NewIfBlock = ThenTerm->getSuccessor(Idx: 0);
680 NewIfBlock->setName("else");
681 BasicBlock *PrevIfBlock = IfBlock;
682 IfBlock = NewIfBlock;
683
684 // Create the phi to join the new and previous value.
685 Builder.SetInsertPoint(TheBB: NewIfBlock, IP: NewIfBlock->begin());
686 PHINode *Phi = Builder.CreatePHI(Ty: VecType, NumReservedValues: 2, Name: "res.phi.else");
687 Phi->addIncoming(V: NewVResult, BB: CondBlock);
688 Phi->addIncoming(V: VResult, BB: PrevIfBlock);
689 VResult = Phi;
690 }
691
692 CI->replaceAllUsesWith(V: VResult);
693 CI->eraseFromParent();
694
695 ModifiedDT = true;
696}
697
698// Translate a masked scatter intrinsic, like
699// void @llvm.masked.scatter.v16i32(<16 x i32> %Src, <16 x i32*>* %Ptrs, i32 4,
700// <16 x i1> %Mask)
701// to a chain of basic blocks, that stores element one-by-one if
702// the appropriate mask bit is set.
703//
704// %Ptrs = getelementptr i32, i32* %ptr, <16 x i64> %ind
705// %Mask0 = extractelement <16 x i1> %Mask, i32 0
706// br i1 %Mask0, label %cond.store, label %else
707//
708// cond.store:
709// %Elt0 = extractelement <16 x i32> %Src, i32 0
710// %Ptr0 = extractelement <16 x i32*> %Ptrs, i32 0
711// store i32 %Elt0, i32* %Ptr0, align 4
712// br label %else
713//
714// else:
715// %Mask1 = extractelement <16 x i1> %Mask, i32 1
716// br i1 %Mask1, label %cond.store1, label %else2
717//
718// cond.store1:
719// %Elt1 = extractelement <16 x i32> %Src, i32 1
720// %Ptr1 = extractelement <16 x i32*> %Ptrs, i32 1
721// store i32 %Elt1, i32* %Ptr1, align 4
722// br label %else2
723// . . .
724static void scalarizeMaskedScatter(const DataLayout &DL,
725 bool HasBranchDivergence, CallInst *CI,
726 DomTreeUpdater *DTU, bool &ModifiedDT) {
727 Value *Src = CI->getArgOperand(i: 0);
728 Value *Ptrs = CI->getArgOperand(i: 1);
729 Value *Mask = CI->getArgOperand(i: 2);
730
731 auto *SrcFVTy = cast<FixedVectorType>(Val: Src->getType());
732
733 assert(
734 isa<VectorType>(Ptrs->getType()) &&
735 isa<PointerType>(cast<VectorType>(Ptrs->getType())->getElementType()) &&
736 "Vector of pointers is expected in masked scatter intrinsic");
737
738 IRBuilder<> Builder(CI->getContext());
739 Instruction *InsertPt = CI;
740 Builder.SetInsertPoint(InsertPt);
741 Builder.SetCurrentDebugLocation(CI->getDebugLoc());
742
743 Align AlignVal = CI->getParamAlign(ArgNo: 1).valueOrOne();
744 unsigned VectorWidth = SrcFVTy->getNumElements();
745
746 // Shorten the way if the mask is a vector of constants.
747 if (isConstantIntVector(Mask)) {
748 for (unsigned Idx = 0; Idx < VectorWidth; ++Idx) {
749 if (cast<Constant>(Val: Mask)->getAggregateElement(Elt: Idx)->isNullValue())
750 continue;
751 Value *OneElt =
752 Builder.CreateExtractElement(Vec: Src, Idx, Name: "Elt" + Twine(Idx));
753 Value *Ptr = Builder.CreateExtractElement(Vec: Ptrs, Idx, Name: "Ptr" + Twine(Idx));
754 StoreInst *Store = Builder.CreateAlignedStore(Val: OneElt, Ptr, Align: AlignVal);
755 copyMetadataForScalarizedStore(Dest&: *Store, Source: *CI, DL,
756 /*SourcePtrOperand=*/1,
757 /*ByteOffset=*/0);
758 }
759 CI->eraseFromParent();
760 return;
761 }
762
763 // If the mask is not v1i1, use scalar bit test operations. This generates
764 // better results on X86 at least.
765 Value *SclrMask = nullptr;
766 if (VectorWidth != 1 && !HasBranchDivergence) {
767 Type *SclrMaskTy = Builder.getIntNTy(N: VectorWidth);
768 SclrMask = Builder.CreateBitCast(V: Mask, DestTy: SclrMaskTy, Name: "scalar_mask");
769 }
770
771 for (unsigned Idx = 0; Idx < VectorWidth; ++Idx) {
772 // Fill the "else" block, created in the previous iteration
773 //
774 // %Mask1 = and i16 %scalar_mask, i32 1 << Idx
775 // %cond = icmp ne i16 %mask_1, 0
776 // br i1 %Mask1, label %cond.store, label %else
777 //
778 // On GPUs, use
779 // %cond = extrectelement %mask, Idx
780 // instead
781 Value *Predicate;
782 if (SclrMask != nullptr) {
783 Value *Mask = Builder.getInt(AI: APInt::getOneBitSet(
784 numBits: VectorWidth, BitNo: adjustForEndian(DL, VectorWidth, Idx)));
785 Predicate = Builder.CreateICmpNE(LHS: Builder.CreateAnd(LHS: SclrMask, RHS: Mask),
786 RHS: Builder.getIntN(N: VectorWidth, C: 0));
787 } else {
788 Predicate = Builder.CreateExtractElement(Vec: Mask, Idx, Name: "Mask" + Twine(Idx));
789 }
790
791 // Create "cond" block
792 //
793 // %Elt1 = extractelement <16 x i32> %Src, i32 1
794 // %Ptr1 = extractelement <16 x i32*> %Ptrs, i32 1
795 // %store i32 %Elt1, i32* %Ptr1
796 //
797 // We mark the branch weights as explicitly unknown given they would only
798 // be derivable from the mask which we do not have VP information for.
799 Instruction *ThenTerm =
800 SplitBlockAndInsertIfThen(Cond: Predicate, SplitBefore: InsertPt, /*Unreachable=*/false,
801 BranchWeights: getExplicitlyUnknownBranchWeightsIfProfiled(
802 F&: *CI->getFunction(), DEBUG_TYPE),
803 DTU);
804
805 BasicBlock *CondBlock = ThenTerm->getParent();
806 CondBlock->setName("cond.store");
807
808 Builder.SetInsertPoint(CondBlock->getTerminator());
809 Value *OneElt = Builder.CreateExtractElement(Vec: Src, Idx, Name: "Elt" + Twine(Idx));
810 Value *Ptr = Builder.CreateExtractElement(Vec: Ptrs, Idx, Name: "Ptr" + Twine(Idx));
811 StoreInst *Store = Builder.CreateAlignedStore(Val: OneElt, Ptr, Align: AlignVal);
812 copyMetadataForScalarizedStore(Dest&: *Store, Source: *CI, DL,
813 /*SourcePtrOperand=*/1,
814 /*ByteOffset=*/0);
815
816 // Create "else" block, fill it in the next iteration
817 BasicBlock *NewIfBlock = ThenTerm->getSuccessor(Idx: 0);
818 NewIfBlock->setName("else");
819
820 Builder.SetInsertPoint(TheBB: NewIfBlock, IP: NewIfBlock->begin());
821 }
822 CI->eraseFromParent();
823
824 ModifiedDT = true;
825}
826
827static void scalarizeMaskedExpandLoad(const DataLayout &DL,
828 bool HasBranchDivergence, CallInst *CI,
829 DomTreeUpdater *DTU, bool &ModifiedDT) {
830 Value *Ptr = CI->getArgOperand(i: 0);
831 Value *Mask = CI->getArgOperand(i: 1);
832 Value *PassThru = CI->getArgOperand(i: 2);
833 Align Alignment = CI->getParamAlign(ArgNo: 0).valueOrOne();
834
835 auto *VecType = cast<FixedVectorType>(Val: CI->getType());
836
837 Type *EltTy = VecType->getElementType();
838
839 IRBuilder<> Builder(CI->getContext());
840 Instruction *InsertPt = CI;
841 BasicBlock *IfBlock = CI->getParent();
842
843 Builder.SetInsertPoint(InsertPt);
844 Builder.SetCurrentDebugLocation(CI->getDebugLoc());
845
846 unsigned VectorWidth = VecType->getNumElements();
847
848 // The result vector
849 Value *VResult = PassThru;
850
851 // Adjust alignment for the scalar instruction.
852 const Align AdjustedAlignment =
853 commonAlignment(A: Alignment, Offset: EltTy->getPrimitiveSizeInBits() / 8);
854
855 // Shorten the way if the mask is a vector of constants.
856 // Create a build_vector pattern, with loads/poisons as necessary and then
857 // shuffle blend with the pass through value.
858 if (isConstantIntVector(Mask)) {
859 unsigned MemIndex = 0;
860 VResult = PoisonValue::get(T: VecType);
861 SmallVector<int, 16> ShuffleMask(VectorWidth, PoisonMaskElem);
862 for (unsigned Idx = 0; Idx < VectorWidth; ++Idx) {
863 Value *InsertElt;
864 if (cast<Constant>(Val: Mask)->getAggregateElement(Elt: Idx)->isNullValue()) {
865 InsertElt = PoisonValue::get(T: EltTy);
866 ShuffleMask[Idx] = Idx + VectorWidth;
867 } else {
868 Value *NewPtr =
869 Builder.CreateConstInBoundsGEP1_32(Ty: EltTy, Ptr, Idx0: MemIndex);
870 LoadInst *Load = Builder.CreateAlignedLoad(
871 Ty: EltTy, Ptr: NewPtr, Align: AdjustedAlignment, Name: "Load" + Twine(Idx));
872 copyMetadataForScalarizedLoad(
873 Dest&: *Load, Source: *CI, DL, /*SourcePtrOperand=*/0,
874 ByteOffset: MemIndex * DL.getTypeAllocSize(Ty: EltTy).getFixedValue());
875 InsertElt = Load;
876 ShuffleMask[Idx] = Idx;
877 ++MemIndex;
878 }
879 VResult = Builder.CreateInsertElement(Vec: VResult, NewElt: InsertElt, Idx,
880 Name: "Res" + Twine(Idx));
881 }
882 VResult = Builder.CreateShuffleVector(V1: VResult, V2: PassThru, Mask: ShuffleMask);
883 CI->replaceAllUsesWith(V: VResult);
884 CI->eraseFromParent();
885 return;
886 }
887
888 // If the mask is not v1i1, use scalar bit test operations. This generates
889 // better results on X86 at least. However, don't do this on GPUs or other
890 // machines with branch divergence, as there, each i1 takes up a register.
891 Value *SclrMask = nullptr;
892 if (VectorWidth != 1 && !HasBranchDivergence) {
893 Type *SclrMaskTy = Builder.getIntNTy(N: VectorWidth);
894 SclrMask = Builder.CreateBitCast(V: Mask, DestTy: SclrMaskTy, Name: "scalar_mask");
895 }
896
897 for (unsigned Idx = 0; Idx < VectorWidth; ++Idx) {
898 // Fill the "else" block, created in the previous iteration
899 //
900 // %res.phi.else3 = phi <16 x i32> [ %11, %cond.load1 ], [ %res.phi.else,
901 // %else ] %mask_1 = extractelement <16 x i1> %mask, i32 Idx br i1 %mask_1,
902 // label %cond.load, label %else
903 //
904 // On GPUs, use
905 // %cond = extrectelement %mask, Idx
906 // instead
907
908 Value *Predicate;
909 if (SclrMask != nullptr) {
910 Value *Mask = Builder.getInt(AI: APInt::getOneBitSet(
911 numBits: VectorWidth, BitNo: adjustForEndian(DL, VectorWidth, Idx)));
912 Predicate = Builder.CreateICmpNE(LHS: Builder.CreateAnd(LHS: SclrMask, RHS: Mask),
913 RHS: Builder.getIntN(N: VectorWidth, C: 0));
914 } else {
915 Predicate = Builder.CreateExtractElement(Vec: Mask, Idx, Name: "Mask" + Twine(Idx));
916 }
917
918 // Create "cond" block
919 //
920 // %EltAddr = getelementptr i32* %1, i32 0
921 // %Elt = load i32* %EltAddr
922 // VResult = insertelement <16 x i32> VResult, i32 %Elt, i32 Idx
923 //
924 // We mark the branch weights as explicitly unknown given they would only
925 // be derivable from the mask which we do not have VP information for.
926 Instruction *ThenTerm =
927 SplitBlockAndInsertIfThen(Cond: Predicate, SplitBefore: InsertPt, /*Unreachable=*/false,
928 BranchWeights: getExplicitlyUnknownBranchWeightsIfProfiled(
929 F&: *CI->getFunction(), DEBUG_TYPE),
930 DTU);
931
932 BasicBlock *CondBlock = ThenTerm->getParent();
933 CondBlock->setName("cond.load");
934
935 Builder.SetInsertPoint(CondBlock->getTerminator());
936 LoadInst *Load = Builder.CreateAlignedLoad(Ty: EltTy, Ptr, Align: AdjustedAlignment);
937 copyMetadataForScalarizedLoad(Dest&: *Load, Source: *CI, DL, /*SourcePtrOperand=*/0,
938 ByteOffset: std::nullopt);
939 Value *NewVResult = Builder.CreateInsertElement(Vec: VResult, NewElt: Load, Idx);
940
941 // Move the pointer if there are more blocks to come.
942 Value *NewPtr;
943 if ((Idx + 1) != VectorWidth)
944 NewPtr = Builder.CreateConstInBoundsGEP1_32(Ty: EltTy, Ptr, Idx0: 1);
945
946 // Create "else" block, fill it in the next iteration
947 BasicBlock *NewIfBlock = ThenTerm->getSuccessor(Idx: 0);
948 NewIfBlock->setName("else");
949 BasicBlock *PrevIfBlock = IfBlock;
950 IfBlock = NewIfBlock;
951
952 // Create the phi to join the new and previous value.
953 Builder.SetInsertPoint(TheBB: NewIfBlock, IP: NewIfBlock->begin());
954 PHINode *ResultPhi = Builder.CreatePHI(Ty: VecType, NumReservedValues: 2, Name: "res.phi.else");
955 ResultPhi->addIncoming(V: NewVResult, BB: CondBlock);
956 ResultPhi->addIncoming(V: VResult, BB: PrevIfBlock);
957 VResult = ResultPhi;
958
959 // Add a PHI for the pointer if this isn't the last iteration.
960 if ((Idx + 1) != VectorWidth) {
961 PHINode *PtrPhi = Builder.CreatePHI(Ty: Ptr->getType(), NumReservedValues: 2, Name: "ptr.phi.else");
962 PtrPhi->addIncoming(V: NewPtr, BB: CondBlock);
963 PtrPhi->addIncoming(V: Ptr, BB: PrevIfBlock);
964 Ptr = PtrPhi;
965 }
966 }
967
968 CI->replaceAllUsesWith(V: VResult);
969 CI->eraseFromParent();
970
971 ModifiedDT = true;
972}
973
974static void scalarizeMaskedCompressStore(const DataLayout &DL,
975 bool HasBranchDivergence, CallInst *CI,
976 DomTreeUpdater *DTU,
977 bool &ModifiedDT) {
978 Value *Src = CI->getArgOperand(i: 0);
979 Value *Ptr = CI->getArgOperand(i: 1);
980 Value *Mask = CI->getArgOperand(i: 2);
981 Align Alignment = CI->getParamAlign(ArgNo: 1).valueOrOne();
982
983 auto *VecType = cast<FixedVectorType>(Val: Src->getType());
984
985 IRBuilder<> Builder(CI->getContext());
986 Instruction *InsertPt = CI;
987 BasicBlock *IfBlock = CI->getParent();
988
989 Builder.SetInsertPoint(InsertPt);
990 Builder.SetCurrentDebugLocation(CI->getDebugLoc());
991
992 Type *EltTy = VecType->getElementType();
993
994 // Adjust alignment for the scalar instruction.
995 const Align AdjustedAlignment =
996 commonAlignment(A: Alignment, Offset: EltTy->getPrimitiveSizeInBits() / 8);
997
998 unsigned VectorWidth = VecType->getNumElements();
999
1000 // Shorten the way if the mask is a vector of constants.
1001 if (isConstantIntVector(Mask)) {
1002 unsigned MemIndex = 0;
1003 for (unsigned Idx = 0; Idx < VectorWidth; ++Idx) {
1004 if (cast<Constant>(Val: Mask)->getAggregateElement(Elt: Idx)->isNullValue())
1005 continue;
1006 Value *OneElt =
1007 Builder.CreateExtractElement(Vec: Src, Idx, Name: "Elt" + Twine(Idx));
1008 Value *NewPtr = Builder.CreateConstInBoundsGEP1_32(Ty: EltTy, Ptr, Idx0: MemIndex);
1009 StoreInst *Store =
1010 Builder.CreateAlignedStore(Val: OneElt, Ptr: NewPtr, Align: AdjustedAlignment);
1011 copyMetadataForScalarizedStore(
1012 Dest&: *Store, Source: *CI, DL, /*SourcePtrOperand=*/1,
1013 ByteOffset: MemIndex * DL.getTypeAllocSize(Ty: EltTy).getFixedValue());
1014 ++MemIndex;
1015 }
1016 CI->eraseFromParent();
1017 return;
1018 }
1019
1020 // If the mask is not v1i1, use scalar bit test operations. This generates
1021 // better results on X86 at least. However, don't do this on GPUs or other
1022 // machines with branch divergence, as there, each i1 takes up a register.
1023 Value *SclrMask = nullptr;
1024 if (VectorWidth != 1 && !HasBranchDivergence) {
1025 Type *SclrMaskTy = Builder.getIntNTy(N: VectorWidth);
1026 SclrMask = Builder.CreateBitCast(V: Mask, DestTy: SclrMaskTy, Name: "scalar_mask");
1027 }
1028
1029 for (unsigned Idx = 0; Idx < VectorWidth; ++Idx) {
1030 // Fill the "else" block, created in the previous iteration
1031 //
1032 // %mask_1 = extractelement <16 x i1> %mask, i32 Idx
1033 // br i1 %mask_1, label %cond.store, label %else
1034 //
1035 // On GPUs, use
1036 // %cond = extrectelement %mask, Idx
1037 // instead
1038 Value *Predicate;
1039 if (SclrMask != nullptr) {
1040 Value *Mask = Builder.getInt(AI: APInt::getOneBitSet(
1041 numBits: VectorWidth, BitNo: adjustForEndian(DL, VectorWidth, Idx)));
1042 Predicate = Builder.CreateICmpNE(LHS: Builder.CreateAnd(LHS: SclrMask, RHS: Mask),
1043 RHS: Builder.getIntN(N: VectorWidth, C: 0));
1044 } else {
1045 Predicate = Builder.CreateExtractElement(Vec: Mask, Idx, Name: "Mask" + Twine(Idx));
1046 }
1047
1048 // Create "cond" block
1049 //
1050 // %OneElt = extractelement <16 x i32> %Src, i32 Idx
1051 // %EltAddr = getelementptr i32* %1, i32 0
1052 // %store i32 %OneElt, i32* %EltAddr
1053 //
1054 // We mark the branch weights as explicitly unknown given they would only
1055 // be derivable from the mask which we do not have VP information for.
1056 Instruction *ThenTerm =
1057 SplitBlockAndInsertIfThen(Cond: Predicate, SplitBefore: InsertPt, /*Unreachable=*/false,
1058 BranchWeights: getExplicitlyUnknownBranchWeightsIfProfiled(
1059 F&: *CI->getFunction(), DEBUG_TYPE),
1060 DTU);
1061
1062 BasicBlock *CondBlock = ThenTerm->getParent();
1063 CondBlock->setName("cond.store");
1064
1065 Builder.SetInsertPoint(CondBlock->getTerminator());
1066 Value *OneElt = Builder.CreateExtractElement(Vec: Src, Idx);
1067 StoreInst *Store =
1068 Builder.CreateAlignedStore(Val: OneElt, Ptr, Align: AdjustedAlignment);
1069 copyMetadataForScalarizedStore(Dest&: *Store, Source: *CI, DL, /*SourcePtrOperand=*/1,
1070 ByteOffset: std::nullopt);
1071
1072 // Move the pointer if there are more blocks to come.
1073 Value *NewPtr;
1074 if ((Idx + 1) != VectorWidth)
1075 NewPtr = Builder.CreateConstInBoundsGEP1_32(Ty: EltTy, Ptr, Idx0: 1);
1076
1077 // Create "else" block, fill it in the next iteration
1078 BasicBlock *NewIfBlock = ThenTerm->getSuccessor(Idx: 0);
1079 NewIfBlock->setName("else");
1080 BasicBlock *PrevIfBlock = IfBlock;
1081 IfBlock = NewIfBlock;
1082
1083 Builder.SetInsertPoint(TheBB: NewIfBlock, IP: NewIfBlock->begin());
1084
1085 // Add a PHI for the pointer if this isn't the last iteration.
1086 if ((Idx + 1) != VectorWidth) {
1087 PHINode *PtrPhi = Builder.CreatePHI(Ty: Ptr->getType(), NumReservedValues: 2, Name: "ptr.phi.else");
1088 PtrPhi->addIncoming(V: NewPtr, BB: CondBlock);
1089 PtrPhi->addIncoming(V: Ptr, BB: PrevIfBlock);
1090 Ptr = PtrPhi;
1091 }
1092 }
1093 CI->eraseFromParent();
1094
1095 ModifiedDT = true;
1096}
1097
1098static void scalarizeMaskedVectorHistogram(const DataLayout &DL, CallInst *CI,
1099 DomTreeUpdater *DTU,
1100 bool &ModifiedDT) {
1101 // If we extend histogram to return a result someday (like the updated vector)
1102 // then we'll need to support it here.
1103 assert(CI->getType()->isVoidTy() && "Histogram with non-void return.");
1104 Value *Ptrs = CI->getArgOperand(i: 0);
1105 Value *Inc = CI->getArgOperand(i: 1);
1106 Value *Mask = CI->getArgOperand(i: 2);
1107
1108 auto *AddrType = cast<FixedVectorType>(Val: Ptrs->getType());
1109 Type *EltTy = Inc->getType();
1110
1111 IRBuilder<> Builder(CI->getContext());
1112 Instruction *InsertPt = CI;
1113 Builder.SetInsertPoint(InsertPt);
1114
1115 Builder.SetCurrentDebugLocation(CI->getDebugLoc());
1116
1117 // FIXME: Do we need to add an alignment parameter to the intrinsic?
1118 unsigned VectorWidth = AddrType->getNumElements();
1119 auto CreateHistogramUpdateValue = [&](IntrinsicInst *CI, Value *Load,
1120 Value *Inc) -> Value * {
1121 Value *UpdateOp;
1122 switch (CI->getIntrinsicID()) {
1123 case Intrinsic::experimental_vector_histogram_add:
1124 UpdateOp = Builder.CreateAdd(LHS: Load, RHS: Inc);
1125 break;
1126 case Intrinsic::experimental_vector_histogram_uadd_sat:
1127 UpdateOp =
1128 Builder.CreateIntrinsic(ID: Intrinsic::uadd_sat, OverloadTypes: {EltTy}, Args: {Load, Inc});
1129 break;
1130 case Intrinsic::experimental_vector_histogram_umin:
1131 UpdateOp = Builder.CreateIntrinsic(ID: Intrinsic::umin, OverloadTypes: {EltTy}, Args: {Load, Inc});
1132 break;
1133 case Intrinsic::experimental_vector_histogram_umax:
1134 UpdateOp = Builder.CreateIntrinsic(ID: Intrinsic::umax, OverloadTypes: {EltTy}, Args: {Load, Inc});
1135 break;
1136
1137 default:
1138 llvm_unreachable("Unexpected histogram intrinsic");
1139 }
1140 return UpdateOp;
1141 };
1142
1143 // Shorten the way if the mask is a vector of constants.
1144 if (isConstantIntVector(Mask)) {
1145 for (unsigned Idx = 0; Idx < VectorWidth; ++Idx) {
1146 if (cast<Constant>(Val: Mask)->getAggregateElement(Elt: Idx)->isNullValue())
1147 continue;
1148 Value *Ptr = Builder.CreateExtractElement(Vec: Ptrs, Idx, Name: "Ptr" + Twine(Idx));
1149 LoadInst *Load = Builder.CreateLoad(Ty: EltTy, Ptr, Name: "Load" + Twine(Idx));
1150 copyMetadataForScalarizedLoad(Dest&: *Load, Source: *CI, DL, /*SourcePtrOperand=*/0,
1151 /*ByteOffset=*/0);
1152 Value *Update =
1153 CreateHistogramUpdateValue(cast<IntrinsicInst>(Val: CI), Load, Inc);
1154 StoreInst *Store = Builder.CreateStore(Val: Update, Ptr);
1155 copyMetadataForScalarizedStore(Dest&: *Store, Source: *CI, DL,
1156 /*SourcePtrOperand=*/0,
1157 /*ByteOffset=*/0);
1158 }
1159 CI->eraseFromParent();
1160 return;
1161 }
1162
1163 for (unsigned Idx = 0; Idx < VectorWidth; ++Idx) {
1164 Value *Predicate =
1165 Builder.CreateExtractElement(Vec: Mask, Idx, Name: "Mask" + Twine(Idx));
1166
1167 // We mark the branch weights as explicitly unknown given they would only
1168 // be derivable from the mask which we do not have VP information for.
1169 Instruction *ThenTerm =
1170 SplitBlockAndInsertIfThen(Cond: Predicate, SplitBefore: InsertPt, /*Unreachable=*/false,
1171 BranchWeights: getExplicitlyUnknownBranchWeightsIfProfiled(
1172 F&: *CI->getFunction(), DEBUG_TYPE),
1173 DTU);
1174
1175 BasicBlock *CondBlock = ThenTerm->getParent();
1176 CondBlock->setName("cond.histogram.update");
1177
1178 Builder.SetInsertPoint(CondBlock->getTerminator());
1179 Value *Ptr = Builder.CreateExtractElement(Vec: Ptrs, Idx, Name: "Ptr" + Twine(Idx));
1180 LoadInst *Load = Builder.CreateLoad(Ty: EltTy, Ptr, Name: "Load" + Twine(Idx));
1181 copyMetadataForScalarizedLoad(Dest&: *Load, Source: *CI, DL, /*SourcePtrOperand=*/0,
1182 /*ByteOffset=*/0);
1183 Value *UpdateOp =
1184 CreateHistogramUpdateValue(cast<IntrinsicInst>(Val: CI), Load, Inc);
1185 StoreInst *Store = Builder.CreateStore(Val: UpdateOp, Ptr);
1186 copyMetadataForScalarizedStore(Dest&: *Store, Source: *CI, DL,
1187 /*SourcePtrOperand=*/0,
1188 /*ByteOffset=*/0);
1189
1190 // Create "else" block, fill it in the next iteration
1191 BasicBlock *NewIfBlock = ThenTerm->getSuccessor(Idx: 0);
1192 NewIfBlock->setName("else");
1193 Builder.SetInsertPoint(TheBB: NewIfBlock, IP: NewIfBlock->begin());
1194 }
1195
1196 CI->eraseFromParent();
1197 ModifiedDT = true;
1198}
1199
1200static bool runImpl(Function &F, const TargetTransformInfo &TTI,
1201 DominatorTree *DT) {
1202 std::optional<DomTreeUpdater> DTU;
1203 if (DT)
1204 DTU.emplace(args&: DT, args: DomTreeUpdater::UpdateStrategy::Lazy);
1205
1206 bool EverMadeChange = false;
1207 bool MadeChange = true;
1208 auto &DL = F.getDataLayout();
1209 bool HasBranchDivergence = TTI.hasBranchDivergence(F: &F);
1210 while (MadeChange) {
1211 MadeChange = false;
1212 for (BasicBlock &BB : llvm::make_early_inc_range(Range&: F)) {
1213 bool ModifiedDTOnIteration = false;
1214 MadeChange |= optimizeBlock(BB, ModifiedDT&: ModifiedDTOnIteration, TTI, DL,
1215 HasBranchDivergence, DTU: DTU ? &*DTU : nullptr);
1216
1217 // Restart BB iteration if the dominator tree of the Function was changed
1218 if (ModifiedDTOnIteration)
1219 break;
1220 }
1221
1222 EverMadeChange |= MadeChange;
1223 }
1224 return EverMadeChange;
1225}
1226
1227bool ScalarizeMaskedMemIntrinLegacyPass::runOnFunction(Function &F) {
1228 auto &TTI = getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
1229 DominatorTree *DT = nullptr;
1230 if (auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>())
1231 DT = &DTWP->getDomTree();
1232 return runImpl(F, TTI, DT);
1233}
1234
1235PreservedAnalyses
1236ScalarizeMaskedMemIntrinPass::run(Function &F, FunctionAnalysisManager &AM) {
1237 auto &TTI = AM.getResult<TargetIRAnalysis>(IR&: F);
1238 auto *DT = AM.getCachedResult<DominatorTreeAnalysis>(IR&: F);
1239 if (!runImpl(F, TTI, DT))
1240 return PreservedAnalyses::all();
1241 PreservedAnalyses PA;
1242 PA.preserve<TargetIRAnalysis>();
1243 PA.preserve<DominatorTreeAnalysis>();
1244 return PA;
1245}
1246
1247static bool optimizeBlock(BasicBlock &BB, bool &ModifiedDT,
1248 const TargetTransformInfo &TTI, const DataLayout &DL,
1249 bool HasBranchDivergence, DomTreeUpdater *DTU) {
1250 bool MadeChange = false;
1251
1252 BasicBlock::iterator CurInstIterator = BB.begin();
1253 while (CurInstIterator != BB.end()) {
1254 if (CallInst *CI = dyn_cast<CallInst>(Val: &*CurInstIterator++))
1255 MadeChange |=
1256 optimizeCallInst(CI, ModifiedDT, TTI, DL, HasBranchDivergence, DTU);
1257 if (ModifiedDT)
1258 return true;
1259 }
1260
1261 return MadeChange;
1262}
1263
1264static bool optimizeCallInst(CallInst *CI, bool &ModifiedDT,
1265 const TargetTransformInfo &TTI,
1266 const DataLayout &DL, bool HasBranchDivergence,
1267 DomTreeUpdater *DTU) {
1268 IntrinsicInst *II = dyn_cast<IntrinsicInst>(Val: CI);
1269 if (II) {
1270 // The scalarization code below does not work for scalable vectors.
1271 if (isa<ScalableVectorType>(Val: II->getType()) ||
1272 any_of(Range: II->args(),
1273 P: [](Value *V) { return isa<ScalableVectorType>(Val: V->getType()); }))
1274 return false;
1275 switch (II->getIntrinsicID()) {
1276 default:
1277 break;
1278 case Intrinsic::experimental_vector_histogram_add:
1279 case Intrinsic::experimental_vector_histogram_uadd_sat:
1280 case Intrinsic::experimental_vector_histogram_umin:
1281 case Intrinsic::experimental_vector_histogram_umax:
1282 if (TTI.isLegalMaskedVectorHistogram(AddrType: CI->getArgOperand(i: 0)->getType(),
1283 DataType: CI->getArgOperand(i: 1)->getType()))
1284 return false;
1285 scalarizeMaskedVectorHistogram(DL, CI, DTU, ModifiedDT);
1286 return true;
1287 case Intrinsic::masked_load:
1288 // Scalarize unsupported vector masked load
1289 if (TTI.isLegalMaskedLoad(
1290 DataType: CI->getType(), Alignment: CI->getParamAlign(ArgNo: 0).valueOrOne(),
1291 AddressSpace: cast<PointerType>(Val: CI->getArgOperand(i: 0)->getType())
1292 ->getAddressSpace(),
1293 MaskKind: isConstantIntVector(Mask: CI->getArgOperand(i: 1))
1294 ? TTI::MaskKind::ConstantMask
1295 : TTI::MaskKind::VariableOrConstantMask))
1296 return false;
1297 scalarizeMaskedLoad(DL, HasBranchDivergence, CI, DTU, ModifiedDT);
1298 return true;
1299 case Intrinsic::masked_store:
1300 if (TTI.isLegalMaskedStore(
1301 DataType: CI->getArgOperand(i: 0)->getType(),
1302 Alignment: CI->getParamAlign(ArgNo: 1).valueOrOne(),
1303 AddressSpace: cast<PointerType>(Val: CI->getArgOperand(i: 1)->getType())
1304 ->getAddressSpace(),
1305 MaskKind: isConstantIntVector(Mask: CI->getArgOperand(i: 2))
1306 ? TTI::MaskKind::ConstantMask
1307 : TTI::MaskKind::VariableOrConstantMask))
1308 return false;
1309 scalarizeMaskedStore(DL, HasBranchDivergence, CI, DTU, ModifiedDT);
1310 return true;
1311 case Intrinsic::masked_gather: {
1312 Align Alignment = CI->getParamAlign(ArgNo: 0).valueOrOne();
1313 Type *LoadTy = CI->getType();
1314 if (TTI.isLegalMaskedGather(DataType: LoadTy, Alignment) &&
1315 !TTI.forceScalarizeMaskedGather(Type: cast<VectorType>(Val: LoadTy), Alignment))
1316 return false;
1317 scalarizeMaskedGather(DL, HasBranchDivergence, CI, DTU, ModifiedDT);
1318 return true;
1319 }
1320 case Intrinsic::masked_scatter: {
1321 Align Alignment = CI->getParamAlign(ArgNo: 1).valueOrOne();
1322 Type *StoreTy = CI->getArgOperand(i: 0)->getType();
1323 if (TTI.isLegalMaskedScatter(DataType: StoreTy, Alignment) &&
1324 !TTI.forceScalarizeMaskedScatter(Type: cast<VectorType>(Val: StoreTy),
1325 Alignment))
1326 return false;
1327 scalarizeMaskedScatter(DL, HasBranchDivergence, CI, DTU, ModifiedDT);
1328 return true;
1329 }
1330 case Intrinsic::masked_expandload:
1331 if (TTI.isLegalMaskedExpandLoad(
1332 DataType: CI->getType(),
1333 Alignment: CI->getAttributes().getParamAttrs(ArgNo: 0).getAlignment().valueOrOne()))
1334 return false;
1335 scalarizeMaskedExpandLoad(DL, HasBranchDivergence, CI, DTU, ModifiedDT);
1336 return true;
1337 case Intrinsic::masked_compressstore:
1338 if (TTI.isLegalMaskedCompressStore(
1339 DataType: CI->getArgOperand(i: 0)->getType(),
1340 Alignment: CI->getAttributes().getParamAttrs(ArgNo: 1).getAlignment().valueOrOne()))
1341 return false;
1342 scalarizeMaskedCompressStore(DL, HasBranchDivergence, CI, DTU,
1343 ModifiedDT);
1344 return true;
1345 }
1346 }
1347
1348 return false;
1349}
1350