xref: /llvm-project/llvm/lib/Target/AMDGPU/AMDGPUAtomicOptimizer.cpp (revision c96dfe0d8bfbc3d4e08af33d5036e2453524b97a)
1 //===-- AMDGPUAtomicOptimizer.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 /// \file
10 /// This pass optimizes atomic operations by using a single lane of a wavefront
11 /// to perform the atomic operation, thus reducing contention on that memory
12 /// location.
13 //
14 //===----------------------------------------------------------------------===//
15 
16 #include "AMDGPU.h"
17 #include "GCNSubtarget.h"
18 #include "llvm/Analysis/LegacyDivergenceAnalysis.h"
19 #include "llvm/CodeGen/TargetPassConfig.h"
20 #include "llvm/IR/IRBuilder.h"
21 #include "llvm/IR/InstVisitor.h"
22 #include "llvm/IR/IntrinsicsAMDGPU.h"
23 #include "llvm/InitializePasses.h"
24 #include "llvm/Target/TargetMachine.h"
25 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
26 
27 #define DEBUG_TYPE "amdgpu-atomic-optimizer"
28 
29 using namespace llvm;
30 using namespace llvm::AMDGPU;
31 
32 namespace {
33 
34 struct ReplacementInfo {
35   Instruction *I;
36   AtomicRMWInst::BinOp Op;
37   unsigned ValIdx;
38   bool ValDivergent;
39 };
40 
41 class AMDGPUAtomicOptimizer : public FunctionPass,
42                               public InstVisitor<AMDGPUAtomicOptimizer> {
43 private:
44   SmallVector<ReplacementInfo, 8> ToReplace;
45   const LegacyDivergenceAnalysis *DA;
46   const DataLayout *DL;
47   DominatorTree *DT;
48   const GCNSubtarget *ST;
49   bool IsPixelShader;
50 
51   Value *buildScan(IRBuilder<> &B, AtomicRMWInst::BinOp Op, Value *V,
52                    Value *const Identity) const;
53   Value *buildShiftRight(IRBuilder<> &B, Value *V, Value *const Identity) const;
54   void optimizeAtomic(Instruction &I, AtomicRMWInst::BinOp Op, unsigned ValIdx,
55                       bool ValDivergent) const;
56 
57 public:
58   static char ID;
59 
60   AMDGPUAtomicOptimizer() : FunctionPass(ID) {}
61 
62   bool runOnFunction(Function &F) override;
63 
64   void getAnalysisUsage(AnalysisUsage &AU) const override {
65     AU.addPreserved<DominatorTreeWrapperPass>();
66     AU.addRequired<LegacyDivergenceAnalysis>();
67     AU.addRequired<TargetPassConfig>();
68   }
69 
70   void visitAtomicRMWInst(AtomicRMWInst &I);
71   void visitIntrinsicInst(IntrinsicInst &I);
72 };
73 
74 } // namespace
75 
76 char AMDGPUAtomicOptimizer::ID = 0;
77 
78 char &llvm::AMDGPUAtomicOptimizerID = AMDGPUAtomicOptimizer::ID;
79 
80 bool AMDGPUAtomicOptimizer::runOnFunction(Function &F) {
81   if (skipFunction(F)) {
82     return false;
83   }
84 
85   DA = &getAnalysis<LegacyDivergenceAnalysis>();
86   DL = &F.getParent()->getDataLayout();
87   DominatorTreeWrapperPass *const DTW =
88       getAnalysisIfAvailable<DominatorTreeWrapperPass>();
89   DT = DTW ? &DTW->getDomTree() : nullptr;
90   const TargetPassConfig &TPC = getAnalysis<TargetPassConfig>();
91   const TargetMachine &TM = TPC.getTM<TargetMachine>();
92   ST = &TM.getSubtarget<GCNSubtarget>(F);
93   IsPixelShader = F.getCallingConv() == CallingConv::AMDGPU_PS;
94 
95   visit(F);
96 
97   const bool Changed = !ToReplace.empty();
98 
99   for (ReplacementInfo &Info : ToReplace) {
100     optimizeAtomic(*Info.I, Info.Op, Info.ValIdx, Info.ValDivergent);
101   }
102 
103   ToReplace.clear();
104 
105   return Changed;
106 }
107 
108 void AMDGPUAtomicOptimizer::visitAtomicRMWInst(AtomicRMWInst &I) {
109   // Early exit for unhandled address space atomic instructions.
110   switch (I.getPointerAddressSpace()) {
111   default:
112     return;
113   case AMDGPUAS::GLOBAL_ADDRESS:
114   case AMDGPUAS::LOCAL_ADDRESS:
115     break;
116   }
117 
118   AtomicRMWInst::BinOp Op = I.getOperation();
119 
120   switch (Op) {
121   default:
122     return;
123   case AtomicRMWInst::Add:
124   case AtomicRMWInst::Sub:
125   case AtomicRMWInst::And:
126   case AtomicRMWInst::Or:
127   case AtomicRMWInst::Xor:
128   case AtomicRMWInst::Max:
129   case AtomicRMWInst::Min:
130   case AtomicRMWInst::UMax:
131   case AtomicRMWInst::UMin:
132     break;
133   }
134 
135   const unsigned PtrIdx = 0;
136   const unsigned ValIdx = 1;
137 
138   // If the pointer operand is divergent, then each lane is doing an atomic
139   // operation on a different address, and we cannot optimize that.
140   if (DA->isDivergentUse(&I.getOperandUse(PtrIdx))) {
141     return;
142   }
143 
144   const bool ValDivergent = DA->isDivergentUse(&I.getOperandUse(ValIdx));
145 
146   // If the value operand is divergent, each lane is contributing a different
147   // value to the atomic calculation. We can only optimize divergent values if
148   // we have DPP available on our subtarget, and the atomic operation is 32
149   // bits.
150   if (ValDivergent &&
151       (!ST->hasDPP() || DL->getTypeSizeInBits(I.getType()) != 32)) {
152     return;
153   }
154 
155   // If we get here, we can optimize the atomic using a single wavefront-wide
156   // atomic operation to do the calculation for the entire wavefront, so
157   // remember the instruction so we can come back to it.
158   const ReplacementInfo Info = {&I, Op, ValIdx, ValDivergent};
159 
160   ToReplace.push_back(Info);
161 }
162 
163 void AMDGPUAtomicOptimizer::visitIntrinsicInst(IntrinsicInst &I) {
164   AtomicRMWInst::BinOp Op;
165 
166   switch (I.getIntrinsicID()) {
167   default:
168     return;
169   case Intrinsic::amdgcn_buffer_atomic_add:
170   case Intrinsic::amdgcn_struct_buffer_atomic_add:
171   case Intrinsic::amdgcn_raw_buffer_atomic_add:
172     Op = AtomicRMWInst::Add;
173     break;
174   case Intrinsic::amdgcn_buffer_atomic_sub:
175   case Intrinsic::amdgcn_struct_buffer_atomic_sub:
176   case Intrinsic::amdgcn_raw_buffer_atomic_sub:
177     Op = AtomicRMWInst::Sub;
178     break;
179   case Intrinsic::amdgcn_buffer_atomic_and:
180   case Intrinsic::amdgcn_struct_buffer_atomic_and:
181   case Intrinsic::amdgcn_raw_buffer_atomic_and:
182     Op = AtomicRMWInst::And;
183     break;
184   case Intrinsic::amdgcn_buffer_atomic_or:
185   case Intrinsic::amdgcn_struct_buffer_atomic_or:
186   case Intrinsic::amdgcn_raw_buffer_atomic_or:
187     Op = AtomicRMWInst::Or;
188     break;
189   case Intrinsic::amdgcn_buffer_atomic_xor:
190   case Intrinsic::amdgcn_struct_buffer_atomic_xor:
191   case Intrinsic::amdgcn_raw_buffer_atomic_xor:
192     Op = AtomicRMWInst::Xor;
193     break;
194   case Intrinsic::amdgcn_buffer_atomic_smin:
195   case Intrinsic::amdgcn_struct_buffer_atomic_smin:
196   case Intrinsic::amdgcn_raw_buffer_atomic_smin:
197     Op = AtomicRMWInst::Min;
198     break;
199   case Intrinsic::amdgcn_buffer_atomic_umin:
200   case Intrinsic::amdgcn_struct_buffer_atomic_umin:
201   case Intrinsic::amdgcn_raw_buffer_atomic_umin:
202     Op = AtomicRMWInst::UMin;
203     break;
204   case Intrinsic::amdgcn_buffer_atomic_smax:
205   case Intrinsic::amdgcn_struct_buffer_atomic_smax:
206   case Intrinsic::amdgcn_raw_buffer_atomic_smax:
207     Op = AtomicRMWInst::Max;
208     break;
209   case Intrinsic::amdgcn_buffer_atomic_umax:
210   case Intrinsic::amdgcn_struct_buffer_atomic_umax:
211   case Intrinsic::amdgcn_raw_buffer_atomic_umax:
212     Op = AtomicRMWInst::UMax;
213     break;
214   }
215 
216   const unsigned ValIdx = 0;
217 
218   const bool ValDivergent = DA->isDivergentUse(&I.getOperandUse(ValIdx));
219 
220   // If the value operand is divergent, each lane is contributing a different
221   // value to the atomic calculation. We can only optimize divergent values if
222   // we have DPP available on our subtarget, and the atomic operation is 32
223   // bits.
224   if (ValDivergent &&
225       (!ST->hasDPP() || DL->getTypeSizeInBits(I.getType()) != 32)) {
226     return;
227   }
228 
229   // If any of the other arguments to the intrinsic are divergent, we can't
230   // optimize the operation.
231   for (unsigned Idx = 1; Idx < I.getNumOperands(); Idx++) {
232     if (DA->isDivergentUse(&I.getOperandUse(Idx))) {
233       return;
234     }
235   }
236 
237   // If we get here, we can optimize the atomic using a single wavefront-wide
238   // atomic operation to do the calculation for the entire wavefront, so
239   // remember the instruction so we can come back to it.
240   const ReplacementInfo Info = {&I, Op, ValIdx, ValDivergent};
241 
242   ToReplace.push_back(Info);
243 }
244 
245 // Use the builder to create the non-atomic counterpart of the specified
246 // atomicrmw binary op.
247 static Value *buildNonAtomicBinOp(IRBuilder<> &B, AtomicRMWInst::BinOp Op,
248                                   Value *LHS, Value *RHS) {
249   CmpInst::Predicate Pred;
250 
251   switch (Op) {
252   default:
253     llvm_unreachable("Unhandled atomic op");
254   case AtomicRMWInst::Add:
255     return B.CreateBinOp(Instruction::Add, LHS, RHS);
256   case AtomicRMWInst::Sub:
257     return B.CreateBinOp(Instruction::Sub, LHS, RHS);
258   case AtomicRMWInst::And:
259     return B.CreateBinOp(Instruction::And, LHS, RHS);
260   case AtomicRMWInst::Or:
261     return B.CreateBinOp(Instruction::Or, LHS, RHS);
262   case AtomicRMWInst::Xor:
263     return B.CreateBinOp(Instruction::Xor, LHS, RHS);
264 
265   case AtomicRMWInst::Max:
266     Pred = CmpInst::ICMP_SGT;
267     break;
268   case AtomicRMWInst::Min:
269     Pred = CmpInst::ICMP_SLT;
270     break;
271   case AtomicRMWInst::UMax:
272     Pred = CmpInst::ICMP_UGT;
273     break;
274   case AtomicRMWInst::UMin:
275     Pred = CmpInst::ICMP_ULT;
276     break;
277   }
278   Value *Cond = B.CreateICmp(Pred, LHS, RHS);
279   return B.CreateSelect(Cond, LHS, RHS);
280 }
281 
282 // Use the builder to create an inclusive scan of V across the wavefront, with
283 // all lanes active.
284 Value *AMDGPUAtomicOptimizer::buildScan(IRBuilder<> &B, AtomicRMWInst::BinOp Op,
285                                         Value *V, Value *const Identity) const {
286   Type *const Ty = V->getType();
287   Module *M = B.GetInsertBlock()->getModule();
288   Function *UpdateDPP =
289       Intrinsic::getDeclaration(M, Intrinsic::amdgcn_update_dpp, Ty);
290 
291   for (unsigned Idx = 0; Idx < 4; Idx++) {
292     V = buildNonAtomicBinOp(
293         B, Op, V,
294         B.CreateCall(UpdateDPP,
295                      {Identity, V, B.getInt32(DPP::ROW_SHR0 | 1 << Idx),
296                       B.getInt32(0xf), B.getInt32(0xf), B.getFalse()}));
297   }
298   if (ST->hasDPPBroadcasts()) {
299     // GFX9 has DPP row broadcast operations.
300     V = buildNonAtomicBinOp(
301         B, Op, V,
302         B.CreateCall(UpdateDPP,
303                      {Identity, V, B.getInt32(DPP::BCAST15), B.getInt32(0xa),
304                       B.getInt32(0xf), B.getFalse()}));
305     V = buildNonAtomicBinOp(
306         B, Op, V,
307         B.CreateCall(UpdateDPP,
308                      {Identity, V, B.getInt32(DPP::BCAST31), B.getInt32(0xc),
309                       B.getInt32(0xf), B.getFalse()}));
310   } else {
311     // On GFX10 all DPP operations are confined to a single row. To get cross-
312     // row operations we have to use permlane or readlane.
313 
314     // Combine lane 15 into lanes 16..31 (and, for wave 64, lane 47 into lanes
315     // 48..63).
316     Value *const PermX = B.CreateIntrinsic(
317         Intrinsic::amdgcn_permlanex16, {},
318         {V, V, B.getInt32(-1), B.getInt32(-1), B.getFalse(), B.getFalse()});
319     V = buildNonAtomicBinOp(
320         B, Op, V,
321         B.CreateCall(UpdateDPP,
322                      {Identity, PermX, B.getInt32(DPP::QUAD_PERM_ID),
323                       B.getInt32(0xa), B.getInt32(0xf), B.getFalse()}));
324     if (!ST->isWave32()) {
325       // Combine lane 31 into lanes 32..63.
326       Value *const Lane31 = B.CreateIntrinsic(Intrinsic::amdgcn_readlane, {},
327                                               {V, B.getInt32(31)});
328       V = buildNonAtomicBinOp(
329           B, Op, V,
330           B.CreateCall(UpdateDPP,
331                        {Identity, Lane31, B.getInt32(DPP::QUAD_PERM_ID),
332                         B.getInt32(0xc), B.getInt32(0xf), B.getFalse()}));
333     }
334   }
335   return V;
336 }
337 
338 // Use the builder to create a shift right of V across the wavefront, with all
339 // lanes active, to turn an inclusive scan into an exclusive scan.
340 Value *AMDGPUAtomicOptimizer::buildShiftRight(IRBuilder<> &B, Value *V,
341                                               Value *const Identity) const {
342   Type *const Ty = V->getType();
343   Module *M = B.GetInsertBlock()->getModule();
344   Function *UpdateDPP =
345       Intrinsic::getDeclaration(M, Intrinsic::amdgcn_update_dpp, Ty);
346 
347   if (ST->hasDPPWavefrontShifts()) {
348     // GFX9 has DPP wavefront shift operations.
349     V = B.CreateCall(UpdateDPP,
350                      {Identity, V, B.getInt32(DPP::WAVE_SHR1), B.getInt32(0xf),
351                       B.getInt32(0xf), B.getFalse()});
352   } else {
353     Function *ReadLane =
354         Intrinsic::getDeclaration(M, Intrinsic::amdgcn_readlane, {});
355     Function *WriteLane =
356         Intrinsic::getDeclaration(M, Intrinsic::amdgcn_writelane, {});
357 
358     // On GFX10 all DPP operations are confined to a single row. To get cross-
359     // row operations we have to use permlane or readlane.
360     Value *Old = V;
361     V = B.CreateCall(UpdateDPP,
362                      {Identity, V, B.getInt32(DPP::ROW_SHR0 + 1),
363                       B.getInt32(0xf), B.getInt32(0xf), B.getFalse()});
364 
365     // Copy the old lane 15 to the new lane 16.
366     V = B.CreateCall(WriteLane, {B.CreateCall(ReadLane, {Old, B.getInt32(15)}),
367                                  B.getInt32(16), V});
368 
369     if (!ST->isWave32()) {
370       // Copy the old lane 31 to the new lane 32.
371       V = B.CreateCall(
372           WriteLane,
373           {B.CreateCall(ReadLane, {Old, B.getInt32(31)}), B.getInt32(32), V});
374 
375       // Copy the old lane 47 to the new lane 48.
376       V = B.CreateCall(
377           WriteLane,
378           {B.CreateCall(ReadLane, {Old, B.getInt32(47)}), B.getInt32(48), V});
379     }
380   }
381 
382   return V;
383 }
384 
385 static APInt getIdentityValueForAtomicOp(AtomicRMWInst::BinOp Op,
386                                          unsigned BitWidth) {
387   switch (Op) {
388   default:
389     llvm_unreachable("Unhandled atomic op");
390   case AtomicRMWInst::Add:
391   case AtomicRMWInst::Sub:
392   case AtomicRMWInst::Or:
393   case AtomicRMWInst::Xor:
394   case AtomicRMWInst::UMax:
395     return APInt::getMinValue(BitWidth);
396   case AtomicRMWInst::And:
397   case AtomicRMWInst::UMin:
398     return APInt::getMaxValue(BitWidth);
399   case AtomicRMWInst::Max:
400     return APInt::getSignedMinValue(BitWidth);
401   case AtomicRMWInst::Min:
402     return APInt::getSignedMaxValue(BitWidth);
403   }
404 }
405 
406 static Value *buildMul(IRBuilder<> &B, Value *LHS, Value *RHS) {
407   const ConstantInt *CI = dyn_cast<ConstantInt>(LHS);
408   return (CI && CI->isOne()) ? RHS : B.CreateMul(LHS, RHS);
409 }
410 
411 void AMDGPUAtomicOptimizer::optimizeAtomic(Instruction &I,
412                                            AtomicRMWInst::BinOp Op,
413                                            unsigned ValIdx,
414                                            bool ValDivergent) const {
415   // Start building just before the instruction.
416   IRBuilder<> B(&I);
417 
418   // If we are in a pixel shader, because of how we have to mask out helper
419   // lane invocations, we need to record the entry and exit BB's.
420   BasicBlock *PixelEntryBB = nullptr;
421   BasicBlock *PixelExitBB = nullptr;
422 
423   // If we're optimizing an atomic within a pixel shader, we need to wrap the
424   // entire atomic operation in a helper-lane check. We do not want any helper
425   // lanes that are around only for the purposes of derivatives to take part
426   // in any cross-lane communication, and we use a branch on whether the lane is
427   // live to do this.
428   if (IsPixelShader) {
429     // Record I's original position as the entry block.
430     PixelEntryBB = I.getParent();
431 
432     Value *const Cond = B.CreateIntrinsic(Intrinsic::amdgcn_ps_live, {}, {});
433     Instruction *const NonHelperTerminator =
434         SplitBlockAndInsertIfThen(Cond, &I, false, nullptr, DT, nullptr);
435 
436     // Record I's new position as the exit block.
437     PixelExitBB = I.getParent();
438 
439     I.moveBefore(NonHelperTerminator);
440     B.SetInsertPoint(&I);
441   }
442 
443   Type *const Ty = I.getType();
444   const unsigned TyBitWidth = DL->getTypeSizeInBits(Ty);
445   auto *const VecTy = FixedVectorType::get(B.getInt32Ty(), 2);
446 
447   // This is the value in the atomic operation we need to combine in order to
448   // reduce the number of atomic operations.
449   Value *const V = I.getOperand(ValIdx);
450 
451   // We need to know how many lanes are active within the wavefront, and we do
452   // this by doing a ballot of active lanes.
453   Type *const WaveTy = B.getIntNTy(ST->getWavefrontSize());
454   CallInst *const Ballot =
455       B.CreateIntrinsic(Intrinsic::amdgcn_ballot, WaveTy, B.getTrue());
456 
457   // We need to know how many lanes are active within the wavefront that are
458   // below us. If we counted each lane linearly starting from 0, a lane is
459   // below us only if its associated index was less than ours. We do this by
460   // using the mbcnt intrinsic.
461   Value *Mbcnt;
462   if (ST->isWave32()) {
463     Mbcnt = B.CreateIntrinsic(Intrinsic::amdgcn_mbcnt_lo, {},
464                               {Ballot, B.getInt32(0)});
465   } else {
466     Value *const BitCast = B.CreateBitCast(Ballot, VecTy);
467     Value *const ExtractLo = B.CreateExtractElement(BitCast, B.getInt32(0));
468     Value *const ExtractHi = B.CreateExtractElement(BitCast, B.getInt32(1));
469     Mbcnt = B.CreateIntrinsic(Intrinsic::amdgcn_mbcnt_lo, {},
470                               {ExtractLo, B.getInt32(0)});
471     Mbcnt =
472         B.CreateIntrinsic(Intrinsic::amdgcn_mbcnt_hi, {}, {ExtractHi, Mbcnt});
473   }
474   Mbcnt = B.CreateIntCast(Mbcnt, Ty, false);
475 
476   Value *const Identity = B.getInt(getIdentityValueForAtomicOp(Op, TyBitWidth));
477 
478   Value *ExclScan = nullptr;
479   Value *NewV = nullptr;
480 
481   // If we have a divergent value in each lane, we need to combine the value
482   // using DPP.
483   if (ValDivergent) {
484     // First we need to set all inactive invocations to the identity value, so
485     // that they can correctly contribute to the final result.
486     NewV = B.CreateIntrinsic(Intrinsic::amdgcn_set_inactive, Ty, {V, Identity});
487 
488     const AtomicRMWInst::BinOp ScanOp =
489         Op == AtomicRMWInst::Sub ? AtomicRMWInst::Add : Op;
490     NewV = buildScan(B, ScanOp, NewV, Identity);
491     ExclScan = buildShiftRight(B, NewV, Identity);
492 
493     // Read the value from the last lane, which has accumlated the values of
494     // each active lane in the wavefront. This will be our new value which we
495     // will provide to the atomic operation.
496     Value *const LastLaneIdx = B.getInt32(ST->getWavefrontSize() - 1);
497     if (TyBitWidth == 64) {
498       Value *const ExtractLo = B.CreateTrunc(NewV, B.getInt32Ty());
499       Value *const ExtractHi =
500           B.CreateTrunc(B.CreateLShr(NewV, 32), B.getInt32Ty());
501       CallInst *const ReadLaneLo = B.CreateIntrinsic(
502           Intrinsic::amdgcn_readlane, {}, {ExtractLo, LastLaneIdx});
503       CallInst *const ReadLaneHi = B.CreateIntrinsic(
504           Intrinsic::amdgcn_readlane, {}, {ExtractHi, LastLaneIdx});
505       Value *const PartialInsert = B.CreateInsertElement(
506           UndefValue::get(VecTy), ReadLaneLo, B.getInt32(0));
507       Value *const Insert =
508           B.CreateInsertElement(PartialInsert, ReadLaneHi, B.getInt32(1));
509       NewV = B.CreateBitCast(Insert, Ty);
510     } else if (TyBitWidth == 32) {
511       NewV = B.CreateIntrinsic(Intrinsic::amdgcn_readlane, {},
512                                {NewV, LastLaneIdx});
513     } else {
514       llvm_unreachable("Unhandled atomic bit width");
515     }
516 
517     // Finally mark the readlanes in the WWM section.
518     NewV = B.CreateIntrinsic(Intrinsic::amdgcn_strict_wwm, Ty, NewV);
519   } else {
520     switch (Op) {
521     default:
522       llvm_unreachable("Unhandled atomic op");
523 
524     case AtomicRMWInst::Add:
525     case AtomicRMWInst::Sub: {
526       // The new value we will be contributing to the atomic operation is the
527       // old value times the number of active lanes.
528       Value *const Ctpop = B.CreateIntCast(
529           B.CreateUnaryIntrinsic(Intrinsic::ctpop, Ballot), Ty, false);
530       NewV = buildMul(B, V, Ctpop);
531       break;
532     }
533 
534     case AtomicRMWInst::And:
535     case AtomicRMWInst::Or:
536     case AtomicRMWInst::Max:
537     case AtomicRMWInst::Min:
538     case AtomicRMWInst::UMax:
539     case AtomicRMWInst::UMin:
540       // These operations with a uniform value are idempotent: doing the atomic
541       // operation multiple times has the same effect as doing it once.
542       NewV = V;
543       break;
544 
545     case AtomicRMWInst::Xor:
546       // The new value we will be contributing to the atomic operation is the
547       // old value times the parity of the number of active lanes.
548       Value *const Ctpop = B.CreateIntCast(
549           B.CreateUnaryIntrinsic(Intrinsic::ctpop, Ballot), Ty, false);
550       NewV = buildMul(B, V, B.CreateAnd(Ctpop, 1));
551       break;
552     }
553   }
554 
555   // We only want a single lane to enter our new control flow, and we do this
556   // by checking if there are any active lanes below us. Only one lane will
557   // have 0 active lanes below us, so that will be the only one to progress.
558   Value *const Cond = B.CreateICmpEQ(Mbcnt, B.getIntN(TyBitWidth, 0));
559 
560   // Store I's original basic block before we split the block.
561   BasicBlock *const EntryBB = I.getParent();
562 
563   // We need to introduce some new control flow to force a single lane to be
564   // active. We do this by splitting I's basic block at I, and introducing the
565   // new block such that:
566   // entry --> single_lane -\
567   //       \------------------> exit
568   Instruction *const SingleLaneTerminator =
569       SplitBlockAndInsertIfThen(Cond, &I, false, nullptr, DT, nullptr);
570 
571   // Move the IR builder into single_lane next.
572   B.SetInsertPoint(SingleLaneTerminator);
573 
574   // Clone the original atomic operation into single lane, replacing the
575   // original value with our newly created one.
576   Instruction *const NewI = I.clone();
577   B.Insert(NewI);
578   NewI->setOperand(ValIdx, NewV);
579 
580   // Move the IR builder into exit next, and start inserting just before the
581   // original instruction.
582   B.SetInsertPoint(&I);
583 
584   const bool NeedResult = !I.use_empty();
585   if (NeedResult) {
586     // Create a PHI node to get our new atomic result into the exit block.
587     PHINode *const PHI = B.CreatePHI(Ty, 2);
588     PHI->addIncoming(UndefValue::get(Ty), EntryBB);
589     PHI->addIncoming(NewI, SingleLaneTerminator->getParent());
590 
591     // We need to broadcast the value who was the lowest active lane (the first
592     // lane) to all other lanes in the wavefront. We use an intrinsic for this,
593     // but have to handle 64-bit broadcasts with two calls to this intrinsic.
594     Value *BroadcastI = nullptr;
595 
596     if (TyBitWidth == 64) {
597       Value *const ExtractLo = B.CreateTrunc(PHI, B.getInt32Ty());
598       Value *const ExtractHi =
599           B.CreateTrunc(B.CreateLShr(PHI, 32), B.getInt32Ty());
600       CallInst *const ReadFirstLaneLo =
601           B.CreateIntrinsic(Intrinsic::amdgcn_readfirstlane, {}, ExtractLo);
602       CallInst *const ReadFirstLaneHi =
603           B.CreateIntrinsic(Intrinsic::amdgcn_readfirstlane, {}, ExtractHi);
604       Value *const PartialInsert = B.CreateInsertElement(
605           UndefValue::get(VecTy), ReadFirstLaneLo, B.getInt32(0));
606       Value *const Insert =
607           B.CreateInsertElement(PartialInsert, ReadFirstLaneHi, B.getInt32(1));
608       BroadcastI = B.CreateBitCast(Insert, Ty);
609     } else if (TyBitWidth == 32) {
610 
611       BroadcastI = B.CreateIntrinsic(Intrinsic::amdgcn_readfirstlane, {}, PHI);
612     } else {
613       llvm_unreachable("Unhandled atomic bit width");
614     }
615 
616     // Now that we have the result of our single atomic operation, we need to
617     // get our individual lane's slice into the result. We use the lane offset
618     // we previously calculated combined with the atomic result value we got
619     // from the first lane, to get our lane's index into the atomic result.
620     Value *LaneOffset = nullptr;
621     if (ValDivergent) {
622       LaneOffset =
623           B.CreateIntrinsic(Intrinsic::amdgcn_strict_wwm, Ty, ExclScan);
624     } else {
625       switch (Op) {
626       default:
627         llvm_unreachable("Unhandled atomic op");
628       case AtomicRMWInst::Add:
629       case AtomicRMWInst::Sub:
630         LaneOffset = buildMul(B, V, Mbcnt);
631         break;
632       case AtomicRMWInst::And:
633       case AtomicRMWInst::Or:
634       case AtomicRMWInst::Max:
635       case AtomicRMWInst::Min:
636       case AtomicRMWInst::UMax:
637       case AtomicRMWInst::UMin:
638         LaneOffset = B.CreateSelect(Cond, Identity, V);
639         break;
640       case AtomicRMWInst::Xor:
641         LaneOffset = buildMul(B, V, B.CreateAnd(Mbcnt, 1));
642         break;
643       }
644     }
645     Value *const Result = buildNonAtomicBinOp(B, Op, BroadcastI, LaneOffset);
646 
647     if (IsPixelShader) {
648       // Need a final PHI to reconverge to above the helper lane branch mask.
649       B.SetInsertPoint(PixelExitBB->getFirstNonPHI());
650 
651       PHINode *const PHI = B.CreatePHI(Ty, 2);
652       PHI->addIncoming(UndefValue::get(Ty), PixelEntryBB);
653       PHI->addIncoming(Result, I.getParent());
654       I.replaceAllUsesWith(PHI);
655     } else {
656       // Replace the original atomic instruction with the new one.
657       I.replaceAllUsesWith(Result);
658     }
659   }
660 
661   // And delete the original.
662   I.eraseFromParent();
663 }
664 
665 INITIALIZE_PASS_BEGIN(AMDGPUAtomicOptimizer, DEBUG_TYPE,
666                       "AMDGPU atomic optimizations", false, false)
667 INITIALIZE_PASS_DEPENDENCY(LegacyDivergenceAnalysis)
668 INITIALIZE_PASS_DEPENDENCY(TargetPassConfig)
669 INITIALIZE_PASS_END(AMDGPUAtomicOptimizer, DEBUG_TYPE,
670                     "AMDGPU atomic optimizations", false, false)
671 
672 FunctionPass *llvm::createAMDGPUAtomicOptimizerPass() {
673   return new AMDGPUAtomicOptimizer();
674 }
675