Fix for PR24179 - mem2reg generates incorrect code

[gilr]

Summary

Patch implements the logic described in main comment of promoteSingleBlockAlloca for loads not preceded by a store:

• If there are no stores before the load, it is (still) replaced by an undef (test2 in the added lit test)
• IF there are any stores, but none preceding the load we bail out and let the rest of Mem2Reg's machinery handle this alloca, as the load may be affected by a following store, depending on the CFG (test1 in the added lit test).

[sanjoy]

LGTM with comments addressed.

lib/Transforms/Utils/PromoteMemoryToRegister.cpp

429 ↗

(On Diff #30161)

I'd make this comment more explicit now:

/// alloca is undefined on only some control flow paths.  e.g. code like
/// this is correct in LLVM IR:
///
///  // A is an alloca with no stores so far
///  for (...) {
///    int t = *A;
///    if (!first_iteration)
///      use(t);
///    *A = 42;
///  }
///

476 ↗

(On Diff #30161)

I'd put the StoresByIndex.empty special case here -- that way all of the logic dealing with "cannot find a previous store" lives at one place.

568 ↗

(On Diff #30161)

I'd use && here.

llvm/trunk/lib/Transforms/Utils/PromoteMemoryToRegister.cpp

 }
 
 /// Many allocas are only used within a single basic block.  If this is the
 /// case, avoid traversing the CFG and inserting a lot of potentially useless
 /// PHI nodes by just performing a single linear pass over the basic block
 /// using the Alloca.
 ///
 /// If we cannot promote this alloca (because it is read before it is written),
-/// return true.  This is necessary in cases where, due to control flow, the
-/// alloca is potentially undefined on some control flow paths.  e.g. code like
-/// this is potentially correct:
-///
-///   for (...) { if (c) { A = undef; undef = B; } }
-///
-/// ... so long as A is not used before undef is set.
-static void promoteSingleBlockAlloca(AllocaInst *AI, const AllocaInfo &Info,
+/// return false.  This is necessary in cases where, due to control flow, the
+/// alloca is undefined only on some control flow paths.  e.g. code like
+/// this is correct in LLVM IR:
+///  // A is an alloca with no stores so far
+///  for (...) {
+///    int t = *A;
+///    if (!first_iteration)
+///      use(t);
+///    *A = 42;
+///  }
+static bool promoteSingleBlockAlloca(AllocaInst *AI, const AllocaInfo &Info,
                                      LargeBlockInfo &LBI,
                                      AliasSetTracker *AST) {
   // The trickiest case to handle is when we have large blocks. Because of this,
   // this code is optimized assuming that large blocks happen.  This does not
   // significantly pessimize the small block case.  This uses LargeBlockInfo to
   // make it efficient to get the index of various operations in the block.
 
   // Walk the use-def list of the alloca, getting the locations of all stores.
for (auto UI = AI->user_begin(), E = AI->user_end(); UI != E;) {
     unsigned LoadIdx = LBI.getInstructionIndex(LI);
 
     // Find the nearest store that has a lower index than this load.
     StoresByIndexTy::iterator I =
         std::lower_bound(StoresByIndex.begin(), StoresByIndex.end(),
                          std::make_pair(LoadIdx,
                                         static_cast<StoreInst *>(nullptr)),
                          less_first());
-
-    if (I == StoresByIndex.begin())
-      // If there is no store before this load, the load takes the undef value.
+    if (I == StoresByIndex.begin()) {
+      if (StoresByIndex.empty())
+        // If there are no stores, the load takes the undef value.
         LI->replaceAllUsesWith(UndefValue::get(LI->getType()));
       else
+        // There is no store before this load, bail out (load may be affected
+        // by the following stores - see main comment).
+        return false;
+    }
+    else
       // Otherwise, there was a store before this load, the load takes its value.
       LI->replaceAllUsesWith(std::prev(I)->second->getOperand(0));
 
     if (AST && LI->getType()->isPointerTy())
       AST->deleteValue(LI);
     LI->eraseFromParent();
     LBI.deleteValue(LI);
   }
static bool promoteSingleBlockAlloca(AllocaInst *AI, const AllocaInfo &Info,
 
   // The alloca's debuginfo can be removed as well.
   if (DbgDeclareInst *DDI = Info.DbgDeclare) {
     DDI->eraseFromParent();
     LBI.deleteValue(DDI);
   }
 
   ++NumLocalPromoted;
+  return true;
 }
 
 void PromoteMem2Reg::run() {
   Function &F = *DT.getRoot()->getParent();
 
   if (AST)
     PointerAllocaValues.resize(Allocas.size());
   AllocaDbgDeclares.resize(Allocas.size());
if (Info.DefiningBlocks.size() == 1) {
         RemoveFromAllocasList(AllocaNum);
         ++NumSingleStore;
         continue;
       }
     }
 
     // If the alloca is only read and written in one basic block, just perform a
     // linear sweep over the block to eliminate it.
-    if (Info.OnlyUsedInOneBlock) {
-      promoteSingleBlockAlloca(AI, Info, LBI, AST);
-
+    if (Info.OnlyUsedInOneBlock &&
+        promoteSingleBlockAlloca(AI, Info, LBI, AST)) {
       // The alloca has been processed, move on.
       RemoveFromAllocasList(AllocaNum);
       continue;
     }
 
     // If we haven't computed a numbering for the BB's in the function, do so
     // now.
     if (BBNumbers.empty()) {

llvm/trunk/test/Transforms/Mem2Reg/pr24179.ll

+; RUN: opt -mem2reg < %s -S | FileCheck %s
+
+declare i32 @def(i32)
+declare i1 @use(i32)
+
+; Special case of a single-BB alloca does not apply here since the load
+; is affected by the following store. Expect this case to be identified
+; and a PHI node to be created.
+define void @test1() {
+; CHECK-LABEL: @test1(
+ entry:
+  %t = alloca i32
+  br label %loop
+
+ loop:
+  %v = load i32, i32* %t
+  %c = call i1 @use(i32 %v)
+; CHECK: [[PHI:%.*]] = phi i32 [ undef, %entry ], [ %n, %loop ]
+; CHECK: call i1 @use(i32 [[PHI]])
+  %n = call i32 @def(i32 7)
+  store i32 %n, i32* %t
+  br i1 %c, label %loop, label %exit
+
+ exit:
+  ret void
+}
+
+; Same as above, except there is no following store. The alloca should just be
+; replaced with an undef
+define void @test2() {
+; CHECK-LABEL: @test2(
+ entry:
+  %t = alloca i32
+  br label %loop
+
+ loop:
+  %v = load i32, i32* %t
+  %c = call i1 @use(i32 %v)
+; CHECK: %c = call i1 @use(i32 undef)
+  br i1 %c, label %loop, label %exit
+
+ exit:
+  ret void
+}