[LoopSimplifyCFG] Update MemorySSA in terminator folding. PR39783

[mkazantsev]

Summary

Terminator folding transform lacks MemorySSA update for memory Phis,
while they exist within MemorySSA analysis. They need exactly the same
type of updates as regular Phis. Failing to update them properly ends up
with inconsistent MemorySSA and manifests in various assertion failures.

This patch adds Memory Phi updates to this transform.

Thanks to @jonpa for finding this!

[asbirlea]

Based on the corresponding Phi updates, the MemoryPhi updates look correct.

Nit: any chance we could merge pr39783_1.ll and pr39783_2.ll into a single file and simplify the tests? e.g. can we keep only @1 in the first file and @func76 in the second; clean-up attributes?

Thanks for the patch!

[mkazantsev]

Thanks Alina!

Unfortunately, initial tests are already simplified by bugpoint, and I am not sure if I'll be able to simplify they further manyally. :(
I suggest to merge the fix and make the test cleanup as a follow-up (if it is possible). Is it OK?

[mkazantsev]

Uh, actually I can. I will update the tests and rebase on top.

[mkazantsev]

Rebased on top of simplified tests.

llvm/trunk/lib/Transforms/Scalar/LoopSimplifyCFG.cpp

 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/LoopUtils.h"
 using namespace llvm;
 
 #define DEBUG_TYPE "loop-simplifycfg"
 
 static cl::opt<bool> EnableTermFolding("enable-loop-simplifycfg-term-folding",
-                                       cl::init(false));
+                                       cl::init(true));
 
 STATISTIC(NumTerminatorsFolded,
           "Number of terminators folded to unconditional branches");
 
 /// If \p BB is a switch or a conditional branch, but only one of its successors
 /// can be reached from this block in runtime, return this successor. Otherwise,
 /// return nullptr.
 static BasicBlock *getOnlyLiveSuccessor(BasicBlock *BB) {

 
 /// Helper class that can turn branches and switches with constant conditions
 /// into unconditional branches.
 class ConstantTerminatorFoldingImpl {
 private:
   Loop &L;
   LoopInfo &LI;
   DominatorTree &DT;
+  MemorySSAUpdater *MSSAU;
 
   // Whether or not the current loop will still exist after terminator constant
   // folding will be done. In theory, there are two ways how it can happen:
   // 1. Loop's latch(es) become unreachable from loop header;
   // 2. Loop's header becomes unreachable from method entry.
   // In practice, the second situation is impossible because we only modify the
   // current loop and its preheader and do not affect preheader's reachibility
   // from any other block. So this variable set to true means that loop's latch
for (BasicBlock *BB : FoldCandidates) {
       unsigned TheOnlySuccDuplicates = 0;
       for (auto *Succ : successors(BB))
         if (Succ != TheOnlySucc) {
           DeadSuccessors.insert(Succ);
           // If our successor lies in a different loop, we don't want to remove
           // the one-input Phi because it is a LCSSA Phi.
           bool PreserveLCSSAPhi = !L.contains(Succ);
           Succ->removePredecessor(BB, PreserveLCSSAPhi);
+          if (MSSAU)
+            MSSAU->removeEdge(BB, Succ);
         } else
           ++TheOnlySuccDuplicates;
 
       assert(TheOnlySuccDuplicates > 0 && "Should be!");
       // If TheOnlySucc was BB's successor more than once, after transform it
       // will be its successor only once. Remove redundant inputs from
       // TheOnlySucc's Phis.
       bool PreserveLCSSAPhi = !L.contains(TheOnlySucc);
       for (unsigned Dup = 1; Dup < TheOnlySuccDuplicates; ++Dup)
         TheOnlySucc->removePredecessor(BB, PreserveLCSSAPhi);
+      if (MSSAU && TheOnlySuccDuplicates > 1)
+        MSSAU->removeDuplicatePhiEdgesBetween(BB, TheOnlySucc);
 
       IRBuilder<> Builder(BB->getContext());
       Instruction *Term = BB->getTerminator();
       Builder.SetInsertPoint(Term);
       Builder.CreateBr(TheOnlySucc);
       Term->eraseFromParent();
 
       for (auto *DeadSucc : DeadSuccessors)
         DTU.deleteEdge(BB, DeadSucc);
 
       ++NumTerminatorsFolded;
     }
   }
 
 public:
-  ConstantTerminatorFoldingImpl(Loop &L, LoopInfo &LI, DominatorTree &DT)
-      : L(L), LI(LI), DT(DT) {}
+  ConstantTerminatorFoldingImpl(Loop &L, LoopInfo &LI, DominatorTree &DT,
+                                MemorySSAUpdater *MSSAU)
+      : L(L), LI(LI), DT(DT), MSSAU(MSSAU) {}
   bool run() {
     assert(L.getLoopLatch() && "Should be single latch!");
 
     // Collect all available information about status of blocks after constant
     // folding.
     analyze();
 
     LLVM_DEBUG(dbgs() << "In function " << L.getHeader()->getParent()->getName()

 #endif
 
     return true;
   }
 };
 
 /// Turn branches and switches with known constant conditions into unconditional
 /// branches.
-static bool constantFoldTerminators(Loop &L, DominatorTree &DT, LoopInfo &LI) {
+static bool constantFoldTerminators(Loop &L, DominatorTree &DT, LoopInfo &LI,
+                                    MemorySSAUpdater *MSSAU) {
   if (!EnableTermFolding)
     return false;
 
   // To keep things simple, only process loops with single latch. We
   // canonicalize most loops to this form. We can support multi-latch if needed.
   if (!L.getLoopLatch())
     return false;
 
-  ConstantTerminatorFoldingImpl BranchFolder(L, LI, DT);
+  ConstantTerminatorFoldingImpl BranchFolder(L, LI, DT, MSSAU);
   return BranchFolder.run();
 }
 
 static bool mergeBlocksIntoPredecessors(Loop &L, DominatorTree &DT,
                                         LoopInfo &LI, MemorySSAUpdater *MSSAU) {
   bool Changed = false;
   DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager);
   // Copy blocks into a temporary array to avoid iterator invalidation issues
static bool mergeBlocksIntoPredecessors(Loop &L, DominatorTree &DT,
   return Changed;
 }
 
 static bool simplifyLoopCFG(Loop &L, DominatorTree &DT, LoopInfo &LI,
                             ScalarEvolution &SE, MemorySSAUpdater *MSSAU) {
   bool Changed = false;
 
   // Constant-fold terminators with known constant conditions.
-  Changed |= constantFoldTerminators(L, DT, LI);
+  Changed |= constantFoldTerminators(L, DT, LI, MSSAU);
 
   // Eliminate unconditional branches by merging blocks into their predecessors.
   Changed |= mergeBlocksIntoPredecessors(L, DT, LI, MSSAU);
 
   if (Changed)
     SE.forgetTopmostLoop(&L);
 
   return Changed;

llvm/trunk/test/Transforms/LoopSimplifyCFG/pr39783.ll

-; XFAIL: *
 ; REQUIRES: asserts
 ; RUN: opt -march=z13 -S -loop-simplifycfg -enable-mssa-loop-dependency -enable-loop-simplifycfg-term-folding 2>&1 < %s | FileCheck %s
 target datalayout = "E-m:e-i1:8:16-i8:8:16-i64:64-f128:64-v128:64-a:8:16-n32:64"
 
 @global = external dso_local local_unnamed_addr global i8, align 2
 @global.1 = external dso_local local_unnamed_addr global i32, align 4
 @global.2 = external dso_local local_unnamed_addr global i32, align 4
 @global.3 = external dso_local local_unnamed_addr global i16, align 2