[Atomic-Expand] Run SimplifyCFG from Atomic-Expand on CAS loop blocks.

[pravinjagtap]

Summary

There are potential benefits in simplifying
CFG just after atomic-expand pass since
it changes control flow.

On AMDGPU targets, for global FP atomic
operations, atomic-expand
pass emits CAS loop which is not efficient.

To optimize atomics AMDGPU target runs
AMDGPUAtomicOptimizer just before
atomic-expand pass.
Running AMDGPUAtomic Optimzer and
atomic expand introduces new control flow,
therefore, running CFG Simplification allows
better codegen.

AArch64 deals with this by inserting an extra
simplifyCFG pass run, which seems excessive.

[pravinjagtap]

Want to take initial feedback on this approach.

[foad]

What kind of simplifications does it do on CAS loops? Is there a test that shows the effect?

[arsenm]

Needs test that shows changes. Also ideally would show it obviates the need for aarch64-enable-atomic-cfg-tidy

[pravinjagtap]

Added Floating Point tests to showcase the effect of running simplify CFG

[pravinjagtap]

llvm/test/CodeGen/AMDGPU/global_atomics_scan_fadd.ll

2165–2184 ↗

(On Diff #548983)

I am not sure whether this is what we are expecting. None of the existing test-cases need update for this change. I am struggling to demonstrate the actual benefits of running SimplifyCFG of CAS blocks.

[arsenm]

llvm/test/CodeGen/AMDGPU/global_atomics_scan_fadd.ll

2165–2184 ↗

(On Diff #548983)

This needs some pure IR tests in test/Transforms/AtomicExpand. You could start by hacking out the aarch64 option and see what breaks for potentially interesting cases.

Does this only do anything if the atomic is in more complex control flow? Does it only do anything if the dominator tree is precomputed? Do you see more changes if you force the dominator tree to be required?

llvm/test/CodeGen/AMDGPU/global_atomics_scan_fmin.ll

50 ↗

(On Diff #548983)

Something's not right because all this is doing is starting to preserve a block name. Does this only do anything if the atomic is in more complex control flow? Does it only do anything if the dominator tree is precomputed? Do you see more changes if you force the dominator tree to be required?

[pravinjagtap]

Added test to showcase the benefits of running simplifyCFG from atomic-expand

[pravinjagtap]

llvm/test/Transforms/AtomicExpand/AMDGPU/expand-atomic-simplify-cfg-CAS-block.ll

24 ↗

(On Diff #549821)

Here, we can observe the potential benefits of running simplify CFG. It simplifies the branching.

[pravinjagtap]

llvm/test/Transforms/AtomicExpand/AMDGPU/expand-atomic-simplify-cfg-CAS-block.ll

24 ↗

(On Diff #549821)

Without this it would have been:

; GFX90A:       atomicrmw.start:	
; GFX90A-NEXT:    [[LOADED:%.*]] = phi float [ [[TMP0]], [[IF]] ], [ [[TMP4:%.*]], [[ATOMICRMW_START]] ]	
; GFX90A-NEXT:    [[NEW:%.*]] = fadd float [[LOADED]], [[IN:%.*]]	
; GFX90A-NEXT:    [[TMP1:%.*]] = bitcast float [[NEW]] to i32	
; GFX90A-NEXT:    [[TMP2:%.*]] = bitcast float [[LOADED]] to i32	
; GFX90A-NEXT:    [[TMP3:%.*]] = cmpxchg ptr addrspace(1) [[OUT]], i32 [[TMP2]], i32 [[TMP1]] seq_cst seq_cst, align 4	
; GFX90A-NEXT:    [[SUCCESS:%.*]] = extractvalue { i32, i1 } [[TMP3]], 1	
; GFX90A-NEXT:    [[NEWLOADED:%.*]] = extractvalue { i32, i1 } [[TMP3]], 0	
; GFX90A-NEXT:    [[TMP4]] = bitcast i32 [[NEWLOADED]] to float	
; GFX90A-NEXT:    br i1 [[SUCCESS]], label [[ATOMICRMW_END:%.*]], label [[ATOMICRMW_START]]	
; GFX90A:       atomicrmw.end:	
; GFX90A-NEXT:    br label [[ENDIF:%.*]]	
; GFX90A:       else:	
; GFX90A-NEXT:    [[TMP5:%.*]] = load float, ptr addrspace(1) [[OUT]], align 4	
; GFX90A-NEXT:    br label [[ATOMICRMW_START2:%.*]]	
; GFX90A:       atomicrmw.start2:	

`

[pravinjagtap]

ping.

[arsenm]

What happens if you remove the aarch64 tidy with this?

llvm/test/CodeGen/AMDGPU/global_atomics_scan_fmin.ll

3196 ↗

(On Diff #549821)

Why were these tests deleted?

llvm/test/Transforms/AtomicExpand/AMDGPU/expand-atomic-simplify-cfg-CAS-block.ll

6 ↗

(On Diff #549821)

Can you precommit the test?

[pravinjagtap]

llvm/test/CodeGen/AMDGPU/global_atomics_scan_fmin.ll

3196 ↗

(On Diff #549821)

I was trying to pre-commit these tests through D157712

[pravinjagtap]

addressed review comment

[arsenm]

In D157495#4598420, @pravinjagtap wrote:

addressed review comment

Haven't tried to delete the AArch64 atomic tidy?

llvm/lib/CodeGen/AtomicExpandPass.cpp

79–80

This didn't account for the require and preserve dom tree

116

Do you need to track this or can you just clean each one up as it happens?

360

Why is RequireAndPreserveDomTree a cl:opt?

[pravinjagtap]

llvm/lib/CodeGen/AtomicExpandPass.cpp

116

Do you need to track this or can you just clean each one up as it happens?

I am clearing this vector at the beginning itself in runOnFunction

[pravinjagtap]

In D157495#4598667, @arsenm wrote:

In D157495#4598420, @pravinjagtap wrote:

addressed review comment

Haven't tried to delete the AArch64 atomic tidy?

TBH, I am not sure how to exactly achieve this.

[arsenm]

In D157495#4598715, @pravinjagtap wrote:

In D157495#4598667, @arsenm wrote:

In D157495#4598420, @pravinjagtap wrote:

addressed review comment

Haven't tried to delete the AArch64 atomic tidy?

TBH, I am not sure how to exactly achieve this.

Delete the option and run of the pass and see if it's equivalently effective in the existing tests to this

llvm/lib/CodeGen/AtomicExpandPass.cpp

116

That's not what I meant, I mean you performed the expansion and can immediately simplify the block without recording it and treating it like a separate pass

[pravinjagtap]

llvm/lib/CodeGen/AtomicExpandPass.cpp

116

I think, this is a cleaner way compared to simplifying these basic blocks when created. We need to pass inputs argument required for simplifyCFG API from runOnFunction to all the way inside insertRMWCmpXchgLoop through member functions and few helper functions.

[pravinjagtap]

In D157495#4598719, @arsenm wrote:

In D157495#4598715, @pravinjagtap wrote:

In D157495#4598667, @arsenm wrote:

In D157495#4598420, @pravinjagtap wrote:

addressed review comment

Haven't tried to delete the AArch64 atomic tidy?

TBH, I am not sure how to exactly achieve this.

Delete the option and run of the pass and see if it's equivalently effective in the existing tests to this

You mean instead of

simplifyCFG(BB, *TTI, RequireAndPreserveDomTree ? &DTU : nullptr,
            SimplifyCFGOptions()
                .forwardSwitchCondToPhi(true)
                .convertSwitchRangeToICmp(true)
                .convertSwitchToLookupTable(true)
                .needCanonicalLoops(false)
                .hoistCommonInsts(true)
                .sinkCommonInsts(true));

just call simplifyCFG(BB, TTI) ?

[pravinjagtap]

In D157495#4598794, @pravinjagtap wrote:

In D157495#4598719, @arsenm wrote:

In D157495#4598715, @pravinjagtap wrote:

In D157495#4598667, @arsenm wrote:

In D157495#4598420, @pravinjagtap wrote:

addressed review comment

Haven't tried to delete the AArch64 atomic tidy?

TBH, I am not sure how to exactly achieve this.

Delete the option and run of the pass and see if it's equivalently effective in the existing tests to this

You mean instead of

simplifyCFG(BB, *TTI, RequireAndPreserveDomTree ? &DTU : nullptr,
            SimplifyCFGOptions()
                .forwardSwitchCondToPhi(true)
                .convertSwitchRangeToICmp(true)
                .convertSwitchToLookupTable(true)
                .needCanonicalLoops(false)
                .hoistCommonInsts(true)
                .sinkCommonInsts(true));

just call simplifyCFG(BB, TTI) ?

Output is identical without these options for the test in expand-atomic-simplify-cfg-CAS-block.ll.

[pravinjagtap]

Do you want me to update the patch without AArch64 atomic tidy options ? I think, relying on default options of simplifyCFG is good option here.

llvm/lib/CodeGen/AtomicExpandPass.cpp

360

Why is RequireAndPreserveDomTree a cl:opt?

This is based on usage of simplifyCFG in https://github.com/llvm/llvm-project/blob/851c248dfcdbf52ee88e4643e59453fcc13501d5/llvm/lib/Target/AMDGPU/AMDGPUUnifyDivergentExitNodes.cpp#L185

[pravinjagtap]

Switched to default options of SimplifyCFG instead of AArch64 atomic tidy options

[pravinjagtap]

Code clean up

[rovka]

Could you please rephrase the commit message? It's not clear to me what using a "canonical pass" instead of a simplifyCFG pass means.

llvm/lib/CodeGen/AtomicExpandPass.cpp

1537

Why are we only keeping track of these blocks? There seem to be lots of other places in this file that split blocks and create new ones. Shouldn't we call simplifyCFG for all of them?

[pravinjagtap]

llvm/lib/CodeGen/AtomicExpandPass.cpp

1537

Why are we only keeping track of these blocks? There seem to be lots of other places in this file that split blocks and create new ones. Shouldn't we call simplifyCFG for all of them?

Targets can configure this simplification using separate pass run e.g. Aarch64 is running simplifyCFG after atomic expand pass https://github.com/llvm/llvm-project/blob/57c090b2ea03937e7c6a08a594532788d01bb813/llvm/lib/Target/AArch64/AArch64TargetMachine.cpp#L557

We think, having separate pass run will be expensive, therefore, for AMDGPU usecase, we are interested in running it on local changes done by atomic-expand. Do you think, calling simplifyCFG on entire functions makes much more sense ?

[rovka]

llvm/lib/CodeGen/AtomicExpandPass.cpp

1537

We think, having separate pass run will be expensive, therefore, for AMDGPU usecase, we are interested in running it on local changes done by atomic-expand. Do you think, calling simplifyCFG on entire functions makes much more sense ?

Ok, I don't know how much compile time this will save for AMDGPU, so I'll let the other reviewers comment on whether or not we want to teach this pass to clean up after itself. But if we decide that we do want it to clean up (i.e. run simplifyCFG only on the blocks that it has added), I think it should:

1. be consistent about it. Right now it creates basic blocks in several different places, but with your patch it only cleans up some of them. If there's a good reason for this, it should be documented (at least in the commit message if not in the code). If there isn't, then at least leave some FIXMEs for the other cases, so people don't have to scratch their heads while looking through this code.

2. be an opt-in behaviour, kind of like how the SimplifyCFG pass has all those settings you can fiddle with when adding it. AtomicExpand is used by several different backends, not just AArch64, and several of them add a full SimplifyCFG run after it (Arm, Hexagon). That SimplifyCFG run may serve to clean up both after AtomicExpand, but potentially also other passes that run before or in between, so it might not make sense to remove the SimplifyCFG run for them. In those cases, it will be useless for AtomicExpand to invoke its own piecemeal SimplifyCFG, so they should be able to run the "fast and messy" AtomicExpand if they want to.

That's just my 2 cents, maybe @arsenm or @foad have different opinions.

[arsenm]

llvm/lib/CodeGen/AtomicExpandPass.cpp

1537

I don't think the compile time is uniquely expensive for AMDGPU, but I would assume just calling simplifycfg on modified blocks would be simpler (as atomics are rare) than running the full CFG pass after the fact

I think it's odd for simplifycfg to be in the codegen pipeline, so a more targeted application seems better

[pravinjagtap]

Rebased.

[pravinjagtap]

Added comments that documents the motivation for this change.

[arsenm]

I still want to see the impact of removing the aarch64 pass

[pravinjagtap]

Experiment: Want to understand the impact of removing the aarch64 pass

Expecting 18 tests to fail. They are not auto-generatable.

llvm/include/llvm/CodeGen/AtomicExpandUtils.h

 //===- AtomicExpandUtils.h - Utilities for expanding atomic instructions --===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CODEGEN_ATOMICEXPANDUTILS_H
 #define LLVM_CODEGEN_ATOMICEXPANDUTILS_H
 
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/Support/AtomicOrdering.h"
 
 namespace llvm {
 
 class AtomicRMWInst;
 class Value;
 

 ///     %new_loaded = extractvalue { iN, i1 } %pair, 0
 ///     %success = extractvalue { iN, i1 } %pair, 1
 /// ; End callback produced IR
 ///     br i1 %success, label %atomicrmw.end, label %loop
 /// atomicrmw.end:
 ///     [...]
 ///
 /// Returns true if the containing function was modified.
-bool expandAtomicRMWToCmpXchg(AtomicRMWInst *AI, CreateCmpXchgInstFun CreateCmpXchg);
+bool expandAtomicRMWToCmpXchg(AtomicRMWInst *AI,
+                              CreateCmpXchgInstFun CreateCmpXchg,
+                              SmallVector<BasicBlock *> &CmpXchgLoopBlocks);
 
 } // end namespace llvm
 
 #endif // LLVM_CODEGEN_ATOMICEXPANDUTILS_H

llvm/lib/CodeGen/AtomicExpandPass.cpp

 // include the use of (intrinsic-based) load-linked/store-conditional loops,
 // AtomicCmpXchg, or type coercions.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLFunctionalExtras.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/Analysis/DomTreeUpdater.h"
 #include "llvm/Analysis/InstSimplifyFolder.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/CodeGen/AtomicExpandUtils.h"
 #include "llvm/CodeGen/RuntimeLibcalls.h"
 #include "llvm/CodeGen/TargetLowering.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/CodeGen/ValueTypes.h"

 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/AtomicOrdering.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetMachine.h"
+#include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/LowerAtomic.h"
 #include <cassert>
 #include <cstdint>
 #include <iterator>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "atomic-expand"
 
 namespace {
 
 class AtomicExpand : public FunctionPass {
   const TargetLowering *TLI = nullptr;
   const DataLayout *DL = nullptr;
+  SmallVector<BasicBlock *> CmpXchgLoopBlocks;
 
 public:
   static char ID; // Pass identification, replacement for typeid
 
   AtomicExpand() : FunctionPass(ID) {
     initializeAtomicExpandPass(*PassRegistry::getPassRegistry());
   }
 
   bool runOnFunction(Function &F) override;
 
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<DominatorTreeWrapperPass>();
+    AU.addRequired<TargetTransformInfoWrapperPass>();
+  }
+
 private:
   bool bracketInstWithFences(Instruction *I, AtomicOrdering Order);
   IntegerType *getCorrespondingIntegerType(Type *T, const DataLayout &DL);
   LoadInst *convertAtomicLoadToIntegerType(LoadInst *LI);
   bool tryExpandAtomicLoad(LoadInst *LI);
   bool expandAtomicLoadToLL(LoadInst *LI);
   bool expandAtomicLoadToCmpXchg(LoadInst *LI);
   StoreInst *convertAtomicStoreToIntegerType(StoreInst *SI);
private:
   void expandAtomicRMWToMaskedIntrinsic(AtomicRMWInst *AI);
   void expandAtomicCmpXchgToMaskedIntrinsic(AtomicCmpXchgInst *CI);
 
   AtomicCmpXchgInst *convertCmpXchgToIntegerType(AtomicCmpXchgInst *CI);
   static Value *insertRMWCmpXchgLoop(
       IRBuilderBase &Builder, Type *ResultType, Value *Addr, Align AddrAlign,
       AtomicOrdering MemOpOrder, SyncScope::ID SSID,
       function_ref<Value *(IRBuilderBase &, Value *)> PerformOp,
-      CreateCmpXchgInstFun CreateCmpXchg);
+      CreateCmpXchgInstFun CreateCmpXchg,
+      SmallVector<BasicBlock *> &CmpXchgLoopBlocks);
   bool tryExpandAtomicCmpXchg(AtomicCmpXchgInst *CI);
 
   bool expandAtomicCmpXchg(AtomicCmpXchgInst *CI);
   bool isIdempotentRMW(AtomicRMWInst *RMWI);
   bool simplifyIdempotentRMW(AtomicRMWInst *RMWI);
 
   bool expandAtomicOpToLibcall(Instruction *I, unsigned Size, Align Alignment,
                                Value *PointerOperand, Value *ValueOperand,
                                Value *CASExpected, AtomicOrdering Ordering,
                                AtomicOrdering Ordering2,
                                ArrayRef<RTLIB::Libcall> Libcalls);
   void expandAtomicLoadToLibcall(LoadInst *LI);
   void expandAtomicStoreToLibcall(StoreInst *LI);
   void expandAtomicRMWToLibcall(AtomicRMWInst *I);
   void expandAtomicCASToLibcall(AtomicCmpXchgInst *I);
 
   friend bool
   llvm::expandAtomicRMWToCmpXchg(AtomicRMWInst *AI,
-                                 CreateCmpXchgInstFun CreateCmpXchg);
+                                 CreateCmpXchgInstFun CreateCmpXchg,
+                                 SmallVector<BasicBlock *> &CmpXchgLoopBlocks);
 };
 
 // IRBuilder to be used for replacement atomic instructions.
 struct ReplacementIRBuilder : IRBuilder<InstSimplifyFolder> {
   // Preserves the DebugLoc from I, and preserves still valid metadata.
   explicit ReplacementIRBuilder(Instruction *I, const DataLayout &DL)
       : IRBuilder(I->getContext(), DL) {
     SetInsertPoint(I);
     this->CollectMetadataToCopy(I, {LLVMContext::MD_pcsections});
   }
 };
 
 } // end anonymous namespace
 
 char AtomicExpand::ID = 0;
 
 char &llvm::AtomicExpandID = AtomicExpand::ID;
 
-INITIALIZE_PASS(AtomicExpand, DEBUG_TYPE, "Expand Atomic instructions", false,
-                false)
-
+INITIALIZE_PASS_BEGIN(AtomicExpand, DEBUG_TYPE, "Expand Atomic instructions",
+                      false, false)
+INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
+INITIALIZE_PASS_END(AtomicExpand, DEBUG_TYPE, "Expand Atomic instructions",
+                    false, false)
 FunctionPass *llvm::createAtomicExpandPass() { return new AtomicExpand(); }
 
 // Helper functions to retrieve the size of atomic instructions.
 static unsigned getAtomicOpSize(LoadInst *LI) {
   const DataLayout &DL = LI->getModule()->getDataLayout();
   return DL.getTypeStoreSize(LI->getType());
 }
 
else if (RMWI) {
           MadeChange = true;
         }
 
         MadeChange |= tryExpandAtomicRMW(RMWI);
       }
     } else if (CASI)
       MadeChange |= tryExpandAtomicCmpXchg(CASI);
   }
+
+  DominatorTreeWrapperPass *const DTW =
+      getAnalysisIfAvailable<DominatorTreeWrapperPass>();
+  DomTreeUpdater DTU(DTW ? &DTW->getDomTree() : nullptr,
+                     DomTreeUpdater::UpdateStrategy::Lazy);
+  auto TTI = &getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
+  for (BasicBlock *BB : CmpXchgLoopBlocks) {
+    simplifyCFG(BB, *TTI, RequireAndPreserveDomTree ? &DTU : nullptr,
+                SimplifyCFGOptions().bonusInstThreshold(2));
+  }
+
   return MadeChange;
 }
 
 bool AtomicExpand::bracketInstWithFences(Instruction *I, AtomicOrdering Order) {
   ReplacementIRBuilder Builder(I, *DL);
 
   auto LeadingFence = TLI->emitLeadingFence(Builder, I, Order);
 
if (ValueSize < MinCASSize) {
                           : SSNs[AI->getSyncScopeID()];
       OptimizationRemarkEmitter ORE(AI->getFunction());
       ORE.emit([&]() {
         return OptimizationRemark(DEBUG_TYPE, "Passed", AI)
                << "A compare and swap loop was generated for an atomic "
                << AI->getOperationName(AI->getOperation()) << " operation at "
                << MemScope << " memory scope";
       });
-      expandAtomicRMWToCmpXchg(AI, createCmpXchgInstFun);
+      expandAtomicRMWToCmpXchg(AI, createCmpXchgInstFun, CmpXchgLoopBlocks);
     }
     return true;
   }
   case TargetLoweringBase::AtomicExpansionKind::MaskedIntrinsic: {
     expandAtomicRMWToMaskedIntrinsic(AI);
     return true;
   }
   case TargetLoweringBase::AtomicExpansionKind::BitTestIntrinsic: {
auto PerformPartwordOp = [&](IRBuilderBase &Builder, Value *Loaded) {
     return performMaskedAtomicOp(AI->getOperation(), Builder, Loaded,
                                  ValOperand_Shifted, AI->getValOperand(), PMV);
   };
 
   Value *OldResult;
   if (ExpansionKind == TargetLoweringBase::AtomicExpansionKind::CmpXChg) {
     OldResult = insertRMWCmpXchgLoop(Builder, PMV.WordType, PMV.AlignedAddr,
                                      PMV.AlignedAddrAlignment, MemOpOrder, SSID,
-                                     PerformPartwordOp, createCmpXchgInstFun);
+                                     PerformPartwordOp, createCmpXchgInstFun,
+                                     CmpXchgLoopBlocks);
   } else {
     assert(ExpansionKind == TargetLoweringBase::AtomicExpansionKind::LLSC);
     OldResult = insertRMWLLSCLoop(Builder, PMV.WordType, PMV.AlignedAddr,
                                   PMV.AlignedAddrAlignment, MemOpOrder,
                                   PerformPartwordOp);
   }
 
   Value *FinalOldResult = extractMaskedValue(Builder, OldResult, PMV);
bool AtomicExpand::simplifyIdempotentRMW(AtomicRMWInst *RMWI) {
   }
   return false;
 }
 
 Value *AtomicExpand::insertRMWCmpXchgLoop(
     IRBuilderBase &Builder, Type *ResultTy, Value *Addr, Align AddrAlign,
     AtomicOrdering MemOpOrder, SyncScope::ID SSID,
     function_ref<Value *(IRBuilderBase &, Value *)> PerformOp,
-    CreateCmpXchgInstFun CreateCmpXchg) {
+    CreateCmpXchgInstFun CreateCmpXchg,
+    SmallVector<BasicBlock *> &CmpXchgLoopBlocks) {
   LLVMContext &Ctx = Builder.getContext();
   BasicBlock *BB = Builder.GetInsertBlock();
   Function *F = BB->getParent();
 
   // Given: atomicrmw some_op iN* %addr, iN %incr ordering
   //
   // The standard expansion we produce is:
   //     [...]
   //     %init_loaded = load atomic iN* %addr
   //     br label %loop
   // loop:
   //     %loaded = phi iN [ %init_loaded, %entry ], [ %new_loaded, %loop ]
   //     %new = some_op iN %loaded, %incr
   //     %pair = cmpxchg iN* %addr, iN %loaded, iN %new
   //     %new_loaded = extractvalue { iN, i1 } %pair, 0
   //     %success = extractvalue { iN, i1 } %pair, 1
   //     br i1 %success, label %atomicrmw.end, label %loop
   // atomicrmw.end:
   //     [...]
   BasicBlock *ExitBB =
       BB->splitBasicBlock(Builder.GetInsertPoint(), "atomicrmw.end");
+  CmpXchgLoopBlocks.push_back(ExitBB);
   BasicBlock *LoopBB = BasicBlock::Create(Ctx, "atomicrmw.start", F, ExitBB);
-
+  CmpXchgLoopBlocks.push_back(LoopBB);
   // The split call above "helpfully" added a branch at the end of BB (to the
   // wrong place), but we want a load. It's easiest to just remove
   // the branch entirely.
   std::prev(BB->end())->eraseFromParent();
   Builder.SetInsertPoint(BB);
   LoadInst *InitLoaded = Builder.CreateAlignedLoad(ResultTy, Addr, AddrAlign);
   Builder.CreateBr(LoopBB);
 
case TargetLoweringBase::AtomicExpansionKind::MaskedIntrinsic:
     expandAtomicCmpXchgToMaskedIntrinsic(CI);
     return true;
   case TargetLoweringBase::AtomicExpansionKind::NotAtomic:
     return lowerAtomicCmpXchgInst(CI);
   }
 }
 
 // Note: This function is exposed externally by AtomicExpandUtils.h
-bool llvm::expandAtomicRMWToCmpXchg(AtomicRMWInst *AI,
-                                    CreateCmpXchgInstFun CreateCmpXchg) {
+bool llvm::expandAtomicRMWToCmpXchg(
+    AtomicRMWInst *AI, CreateCmpXchgInstFun CreateCmpXchg,
+    SmallVector<BasicBlock *> &CmpXchgLoopBlocks) {
   ReplacementIRBuilder Builder(AI, AI->getModule()->getDataLayout());
   Builder.setIsFPConstrained(
       AI->getFunction()->hasFnAttribute(Attribute::StrictFP));
 
   // FIXME: If FP exceptions are observable, we should force them off for the
   // loop for the FP atomics.
   Value *Loaded = AtomicExpand::insertRMWCmpXchgLoop(
       Builder, AI->getType(), AI->getPointerOperand(), AI->getAlign(),
       AI->getOrdering(), AI->getSyncScopeID(),
       [&](IRBuilderBase &Builder, Value *Loaded) {
         return buildAtomicRMWValue(AI->getOperation(), Builder, Loaded,
                                    AI->getValOperand());
       },
-      CreateCmpXchg);
+      CreateCmpXchg, CmpXchgLoopBlocks);
 
   AI->replaceAllUsesWith(Loaded);
   AI->eraseFromParent();
   return true;
 }
 
 // In order to use one of the sized library calls such as
 // __atomic_fetch_add_4, the alignment must be sufficient, the size
Success = expandAtomicOpToLibcall(
         nullptr, I->getOrdering(), AtomicOrdering::NotAtomic, Libcalls);
 
   // The expansion failed: either there were no libcalls at all for
   // the operation (min/max), or there were only size-specialized
   // libcalls (add/sub/etc) and we needed a generic. So, expand to a
   // CAS libcall, via a CAS loop, instead.
   if (!Success) {
     expandAtomicRMWToCmpXchg(
-        I, [this](IRBuilderBase &Builder, Value *Addr, Value *Loaded,
-                  Value *NewVal, Align Alignment, AtomicOrdering MemOpOrder,
-                  SyncScope::ID SSID, Value *&Success, Value *&NewLoaded) {
+        I,
+        [this](IRBuilderBase &Builder, Value *Addr, Value *Loaded,
           // Create the CAS instruction normally...
           AtomicCmpXchgInst *Pair = Builder.CreateAtomicCmpXchg(
               Addr, Loaded, NewVal, Alignment, MemOpOrder,
               AtomicCmpXchgInst::getStrongestFailureOrdering(MemOpOrder), SSID);
           Success = Builder.CreateExtractValue(Pair, 1, "success");
           NewLoaded = Builder.CreateExtractValue(Pair, 0, "newloaded");
 
           // ...and then expand the CAS into a libcall.
           expandAtomicCASToLibcall(Pair);
-        });
+        },
+        CmpXchgLoopBlocks);
   }
 }
 
 // A helper routine for the above expandAtomic*ToLibcall functions.
 //
 // 'Libcalls' contains an array of enum values for the particular
 // ATOMIC libcalls to be emitted. All of the other arguments besides
 // 'I' are extracted from the Instruction subclass by the