[SCEVExpander] Move LCSSA fixup to ::expand.

[fhahn]

Summary

Move LCSSA fixup from ::expandCodeForImpl to ::expand(). This has
the advantage that we directly preserve LCSSA nodes here instead of
relying on doing so in rememberInstruction. It also ensures that we
don't add the non-LCSSA-safe value to InsertedExpressions.

Alternative to D132704.

Fixes #57000.

[efriedma]

Maybe we should restructure the InsertedExpressions cache so it's more likely to consistently hit?

Insertion points fall into the following categories:

• A passed-in arbitrary point.
• The beginning of a basic block.
• The end of a basic block.

The values are then adjusted in some cases: if the insertion point points at an inserted instruction, we adjust forward to try to reuse those instructions.

For the "beginning of a basic block" case, it doesn't really make sense to track the exact instruction where we can actually insert instructions; all that matters is that it's the "beginning". If we fix that, it should ensure we hit the cache reliably for those cases.

Doing multiple cache lookups like the current patch may sort of work in this case, but it's not clear in general why we would expect the unadjusted insertion point to be in the cache.

[fhahn]

In D134739#3819071, @efriedma wrote:

Maybe we should restructure the InsertedExpressions cache so it's more likely to consistently hit?

Insertion points fall into the following categories:

• A passed-in arbitrary point.
• The beginning of a basic block.
• The end of a basic block.

The values are then adjusted in some cases: if the insertion point points at an inserted instruction, we adjust forward to try to reuse those instructions.

For the "beginning of a basic block" case, it doesn't really make sense to track the exact instruction where we can actually insert instructions; all that matters is that it's the "beginning". If we fix that, it should ensure we hit the cache reliably for those cases.

Doing multiple cache lookups like the current patch may sort of work in this case, but it's not clear in general why we would expect the unadjusted insertion point to be in the cache.

I took another closer look at where the LCSSA phi nodes should be cached. I restructured the code to move LCSSA fixup to ::expand() from ::expandCodeForImpl. This has the advantage that we directly preserve LCSSA nodes here instead of relying on doing so in rememberInstruction. It also ensures that we don't add the non-LCSSA-safe value to InsertedExpressions.

Overall the code should be mode direct, this leads to better results in a few cases (fewer unnecessary LCSSA phis) and fixes the crash in the newly added test case.

What do you think?

[efriedma]

LGTM with one minor comment

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

2510

"SH" is unused?

[fhahn]

Remove unused argument, thanks!

[fhahn]

Thanks for the review Eli!

llvm/include/llvm/Transforms/Utils/ScalarEvolutionExpander.h

private:
   PHINode *getAddRecExprPHILiterally(const SCEVAddRecExpr *Normalized,
                                      const Loop *L, Type *ExpandTy, Type *IntTy,
                                      Type *&TruncTy, bool &InvertStep);
   Value *expandIVInc(PHINode *PN, Value *StepV, const Loop *L, Type *ExpandTy,
                      Type *IntTy, bool useSubtract);
 
   void fixupInsertPoints(Instruction *I);
 
-  /// If required, create LCSSA PHIs for \p Users' operand \p OpIdx. If new
-  /// LCSSA PHIs have been created, return the LCSSA PHI available at \p User.
-  /// If no PHIs have been created, return the unchanged operand \p OpIdx.
-  Value *fixupLCSSAFormFor(Instruction *User, unsigned OpIdx);
+  /// Create LCSSA PHIs for \p V, if it is required for uses at the Builder's
+  /// current insertion point.
+  Value *fixupLCSSAFormFor(Value *V, const SCEV *SH);
 };
 
 /// Helper to remove instructions inserted during SCEV expansion, unless they
 /// are marked as used.
 class SCEVExpanderCleaner {
   SCEVExpander &Expander;
 
   /// Indicates whether the result of the expansion is used. If false, the

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

 //===- ScalarEvolutionExpander.cpp - Scalar Evolution Analysis ------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains the implementation of the scalar evolution expander,
 // which is used to generate the code corresponding to a given scalar evolution
 // expression.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Utils/ScalarEvolutionExpander.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/IntrinsicInst.h"
Value *SCEVExpander::expandCodeForImpl(const SCEV *SH, Type *Ty,
   Value *V = expandCodeForImpl(SH, Ty);
   return V;
 }
 
 Value *SCEVExpander::expandCodeForImpl(const SCEV *SH, Type *Ty) {
   // Expand the code for this SCEV.
   Value *V = expand(SH);
 
-  if (PreserveLCSSA) {
-    if (auto *Inst = dyn_cast<Instruction>(V)) {
-      // Create a temporary instruction to at the current insertion point, so we
-      // can hand it off to the helper to create LCSSA PHIs if required for the
-      // new use.
-      // FIXME: Ideally formLCSSAForInstructions (used in fixupLCSSAFormFor)
-      // would accept a insertion point and return an LCSSA phi for that
-      // insertion point, so there is no need to insert & remove the temporary
-      // instruction.
-      Type *ToTy;
-      if (Inst->getType()->isIntegerTy())
-        ToTy = Inst->getType()->getPointerTo();
-      else
-        ToTy = Type::getInt32Ty(Inst->getContext());
-      Instruction *Tmp = CastInst::CreateBitOrPointerCast(
-          Inst, ToTy, "tmp.lcssa.user", &*Builder.GetInsertPoint());
-      V = fixupLCSSAFormFor(Tmp, 0);
-
-      // Clean up temporary instruction.
-      Tmp->eraseFromParent();
-    }
-  }
-
-  InsertedExpressions[std::make_pair(SH, &*Builder.GetInsertPoint())] = V;
   if (Ty) {
     assert(SE.getTypeSizeInBits(Ty) == SE.getTypeSizeInBits(SH->getType()) &&
            "non-trivial casts should be done with the SCEVs directly!");
     V = InsertNoopCastOfTo(V, Ty);
   }
   return V;
 }
 
Value *SCEVExpander::expand(const SCEV *S) {
   if (I != InsertedExpressions.end())
     return I->second;
 
   SCEVInsertPointGuard Guard(Builder, this);
   Builder.SetInsertPoint(InsertPt);
 
   // Expand the expression into instructions.
   Value *V = FindValueInExprValueMap(S, InsertPt);
-  if (!V)
+  if (!V) {
     V = visit(S);
-  else {
+    V = fixupLCSSAFormFor(V, S);
+  } else {
     // If we're reusing an existing instruction, we are effectively CSEing two
     // copies of the instruction (with potentially different flags).  As such,
     // we need to drop any poison generating flags unless we can prove that
     // said flags must be valid for all new users.
     if (auto *I = dyn_cast<Instruction>(V))
       if (I->hasPoisonGeneratingFlags() && !programUndefinedIfPoison(I))
         I->dropPoisonGeneratingFlags();
   }

 void SCEVExpander::rememberInstruction(Value *I) {
   auto DoInsert = [this](Value *V) {
     if (!PostIncLoops.empty())
       InsertedPostIncValues.insert(V);
     else
       InsertedValues.insert(V);
   };
   DoInsert(I);
-
-  if (!PreserveLCSSA)
-    return;
-
-  if (auto *Inst = dyn_cast<Instruction>(I)) {
-    // A new instruction has been added, which might introduce new uses outside
-    // a defining loop. Fix LCSSA from for each operand of the new instruction,
-    // if required.
-    for (unsigned OpIdx = 0, OpEnd = Inst->getNumOperands(); OpIdx != OpEnd;
-         OpIdx++)
-      fixupLCSSAFormFor(Inst, OpIdx);
-  }
 }
 
 /// replaceCongruentIVs - Check for congruent phis in this loop header and
 /// replace them with their most canonical representative. Return the number of
 /// phis eliminated.
 ///
 /// This does not depend on any SCEVExpander state but should be used in
 /// the same context that SCEVExpander is used.
for (const auto *Pred : Union->getPredicates()) {
     Builder.SetInsertPoint(IP);
   }
 
   if (Checks.empty())
     return ConstantInt::getFalse(IP->getContext());
   return Builder.CreateOr(Checks);
 }
 
-Value *SCEVExpander::fixupLCSSAFormFor(Instruction *User, unsigned OpIdx) {
-  assert(PreserveLCSSA);
-  SmallVector<Instruction *, 1> ToUpdate;
-
-  auto *OpV = User->getOperand(OpIdx);
-  auto *OpI = dyn_cast<Instruction>(OpV);
-  if (!OpI)
-    return OpV;
+Value *SCEVExpander::fixupLCSSAFormFor(Value *V, const SCEV *SH) {
+  auto *DefI = dyn_cast<Instruction>(V);
+  if (!PreserveLCSSA || !DefI)
+    return V;
 
-  Loop *DefLoop = SE.LI.getLoopFor(OpI->getParent());
-  Loop *UseLoop = SE.LI.getLoopFor(User->getParent());
+  Instruction *InsertPt = &*Builder.GetInsertPoint();
+  Loop *DefLoop = SE.LI.getLoopFor(DefI->getParent());
+  Loop *UseLoop = SE.LI.getLoopFor(InsertPt->getParent());
   if (!DefLoop || UseLoop == DefLoop || DefLoop->contains(UseLoop))
-    return OpV;
+    return V;
 
-  ToUpdate.push_back(OpI);
+  // Create a temporary instruction to at the current insertion point, so we
+  // can hand it off to the helper to create LCSSA PHIs if required for the
+  // new use.
+  // FIXME: Ideally formLCSSAForInstructions (used in fixupLCSSAFormFor)
+  // would accept a insertion point and return an LCSSA phi for that
+  // insertion point, so there is no need to insert & remove the temporary
+  // instruction.
+  Type *ToTy;
+  if (DefI->getType()->isIntegerTy())
+    ToTy = DefI->getType()->getPointerTo();
+  else
+    ToTy = Type::getInt32Ty(DefI->getContext());
+  Instruction *User =
+      CastInst::CreateBitOrPointerCast(DefI, ToTy, "tmp.lcssa.user", InsertPt);
+  auto RemoveUserOnExit =
+      make_scope_exit([User]() { User->eraseFromParent(); });
+
+  SmallVector<Instruction *, 1> ToUpdate;
+  ToUpdate.push_back(DefI);
   SmallVector<PHINode *, 16> PHIsToRemove;
   formLCSSAForInstructions(ToUpdate, SE.DT, SE.LI, &SE, Builder, &PHIsToRemove);
   for (PHINode *PN : PHIsToRemove) {
     if (!PN->use_empty())
       continue;
     InsertedValues.erase(PN);
     InsertedPostIncValues.erase(PN);
     PN->eraseFromParent();
   }
 
-  return User->getOperand(OpIdx);
+  return User->getOperand(0);
 }
 
 namespace {
 // Search for a SCEV subexpression that is not safe to expand.  Any expression
 // that may expand to a !isSafeToSpeculativelyExecute value is unsafe, namely
 // UDiv expressions. We don't know if the UDiv is derived from an IR divide
 // instruction, but the important thing is that we prove the denominator is
 // nonzero before expansion.

llvm/test/Transforms/IndVarSimplify/lcssa-preservation.ll

 loop1.latch:
   %res.lcssa = phi i64 [ %res, %loop2 ]
   br i1 false, label %loop1, label %exit
 
 exit:
   %res.lcssa2 = phi i64 [ %res.lcssa, %loop1.latch ]
   ret i64 %res.lcssa2
 }
+
+; Check that it does not crash because the incoming values of an LCSSA phi are
+; equal.
+define void @pr57000(i64 %a) {
+; CHECK-LABEL: @pr57000(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP_1_OUTER:%.*]]
+; CHECK:       loop.1.loopexit:
+; CHECK-NEXT:    br label [[LOOP_1_OUTER]]
+; CHECK:       loop.1.outer:
+; CHECK-NEXT:    br label [[LOOP_1:%.*]]
+; CHECK:       loop.1:
+; CHECK-NEXT:    [[CMP:%.*]] = icmp sle i64 [[A:%.*]], 2546175499358690212
+; CHECK-NEXT:    [[CMP_EXT:%.*]] = zext i1 [[CMP]] to i8
+; CHECK-NEXT:    br i1 [[CMP]], label [[LOOP_1]], label [[LOOP_2_HEADER_PREHEADER:%.*]]
+; CHECK:       loop.2.header.preheader:
+; CHECK-NEXT:    [[CMP_LCSSA:%.*]] = phi i1 [ [[CMP]], [[LOOP_1]] ]
+; CHECK-NEXT:    [[CMP_EXT_LCSSA:%.*]] = phi i8 [ [[CMP_EXT]], [[LOOP_1]] ]
+; CHECK-NEXT:    br label [[LOOP_2_HEADER_OUTER:%.*]]
+; CHECK:       loop.2.header.outer:
+; CHECK-NEXT:    br label [[LOOP_2_HEADER:%.*]]
+; CHECK:       loop.2.header:
+; CHECK-NEXT:    switch i8 [[CMP_EXT_LCSSA]], label [[LOOP_1_LOOPEXIT:%.*]] [
+; CHECK-NEXT:    i8 -1, label [[LOOP_2_LATCH:%.*]]
+; CHECK-NEXT:    i8 1, label [[LOOP_2_LATCH]]
+; CHECK-NEXT:    i8 4, label [[LOOP_2_HEADER]]
+; CHECK-NEXT:    ]
+; CHECK:       loop.2.latch:
+; CHECK-NEXT:    [[CMP_TRUNC_LCSSA1:%.*]] = phi i1 [ [[CMP_LCSSA]], [[LOOP_2_HEADER]] ], [ [[CMP_LCSSA]], [[LOOP_2_HEADER]] ]
+; CHECK-NEXT:    call void @use(i1 [[CMP_TRUNC_LCSSA1]])
+; CHECK-NEXT:    br label [[LOOP_2_HEADER_OUTER]]
+;
+entry:
+  br label %loop.1
+
+loop.1:
+  %p.1 = phi i1 [ 0 , %entry ], [ %p.1, %loop.1 ], [ %p.2, %loop.2.header ]
+  %cmp = icmp sle i64 %a, 2546175499358690212
+  %cmp.ext = zext i1 %cmp to i8
+  br i1 %cmp, label %loop.1, label %loop.2.header
+
+loop.2.header:
+  %p.2 = phi i1 [ %p.1, %loop.1 ], [ %p.2, %loop.2.header ], [ %cmp, %loop.2.latch ]
+  %cmp.trunc = trunc i8 %cmp.ext to i1
+  switch i8 %cmp.ext, label %loop.1 [
+  i8 -1, label %loop.2.latch
+  i8 1, label %loop.2.latch
+  i8 4, label %loop.2.header
+  ]
+
+loop.2.latch:
+  call void @use(i1 %cmp.trunc)
+  br label %loop.2.header
+}
+
+declare void @use(i1)

llvm/test/Transforms/LoopVectorize/skeleton-lcssa-crash.ll

 ; CHECK-NEXT:    br label [[INNER:%.*]]
 ; CHECK:       inner:
 ; CHECK-NEXT:    [[C_2:%.*]] = call i1 @cond()
 ; CHECK-NEXT:    br i1 [[C_2]], label [[OUTER_LATCH:%.*]], label [[INNER_BB:%.*]]
 ; CHECK:       inner.bb:
 ; CHECK-NEXT:    [[C_3:%.*]] = call i1 @cond()
 ; CHECK-NEXT:    br i1 [[C_3]], label [[LOOP_3_PREHEADER:%.*]], label [[INNER_LATCH:%.*]]
 ; CHECK:       loop.3.preheader:
-; CHECK-NEXT:    [[L_1_LCSSA7:%.*]] = phi ptr [ [[L_1]], [[INNER_BB]] ]
 ; CHECK-NEXT:    [[L_1_LCSSA:%.*]] = phi ptr [ [[L_1]], [[INNER_BB]] ]
 ; CHECK-NEXT:    [[L_2_LCSSA:%.*]] = phi ptr [ [[L_2]], [[INNER_BB]] ]
 ; CHECK-NEXT:    [[TMP0:%.*]] = add i64 [[N:%.*]], 1
 ; CHECK-NEXT:    [[MIN_ITERS_CHECK:%.*]] = icmp ult i64 [[TMP0]], 2
 ; CHECK-NEXT:    br i1 [[MIN_ITERS_CHECK]], label [[SCALAR_PH:%.*]], label [[VECTOR_MEMCHECK:%.*]]
 ; CHECK:       vector.memcheck:
 ; CHECK-NEXT:    [[UGLYGEP:%.*]] = getelementptr i8, ptr [[L_2_LCSSA]], i64 2
-; CHECK-NEXT:    [[UGLYGEP3:%.*]] = getelementptr i8, ptr [[L_1_LCSSA]], i64 2
+; CHECK-NEXT:    [[UGLYGEP5:%.*]] = getelementptr i8, ptr [[L_1_LCSSA]], i64 2
 ; CHECK-NEXT:    [[TMP1:%.*]] = shl i64 [[N]], 1
 ; CHECK-NEXT:    [[TMP2:%.*]] = add i64 [[TMP1]], 4
-; CHECK-NEXT:    [[UGLYGEP6:%.*]] = getelementptr i8, ptr [[L_1_LCSSA7]], i64 [[TMP2]]
+; CHECK-NEXT:    [[UGLYGEP6:%.*]] = getelementptr i8, ptr [[L_1_LCSSA]], i64 [[TMP2]]
 ; CHECK-NEXT:    [[BOUND0:%.*]] = icmp ult ptr [[L_2_LCSSA]], [[UGLYGEP6]]
-; CHECK-NEXT:    [[BOUND1:%.*]] = icmp ult ptr [[UGLYGEP3]], [[UGLYGEP]]
+; CHECK-NEXT:    [[BOUND1:%.*]] = icmp ult ptr [[UGLYGEP5]], [[UGLYGEP]]
 ; CHECK-NEXT:    [[FOUND_CONFLICT:%.*]] = and i1 [[BOUND0]], [[BOUND1]]
 ; CHECK-NEXT:    br i1 [[FOUND_CONFLICT]], label [[SCALAR_PH]], label [[VECTOR_PH:%.*]]
 ; CHECK:       vector.ph:
 ; CHECK-NEXT:    [[N_MOD_VF:%.*]] = urem i64 [[TMP0]], 2
 ; CHECK-NEXT:    [[N_VEC:%.*]] = sub i64 [[TMP0]], [[N_MOD_VF]]
 ; CHECK-NEXT:    br label [[VECTOR_BODY:%.*]]
 ; CHECK:       vector.body:
 ; CHECK-NEXT:    [[INDEX:%.*]] = phi i64 [ 0, [[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], [[VECTOR_BODY]] ]

 ; CHECK-NEXT:    [[GEP_1:%.*]] = getelementptr inbounds i16, ptr [[L_1_LCSSA]], i64 [[IV_NEXT]]
 ; CHECK-NEXT:    [[LOOP_L_1:%.*]] = load i16, ptr [[GEP_1]], align 2
 ; CHECK-NEXT:    [[GEP_2:%.*]] = getelementptr inbounds i16, ptr [[L_2_LCSSA]], i64 0
 ; CHECK-NEXT:    store i16 [[LOOP_L_1]], ptr [[GEP_2]], align 2
 ; CHECK-NEXT:    br i1 [[C_5]], label [[LOOP_3]], label [[EXIT_LOOPEXIT]], !llvm.loop [[LOOP7:![0-9]+]]
 ; CHECK:       exit.loopexit:
 ; CHECK-NEXT:    br label [[EXIT:%.*]]
 ; CHECK:       exit.loopexit1:
-; CHECK-NEXT:    [[L_1_LCSSA4:%.*]] = phi ptr [ [[L_1]], [[INNER_LATCH]] ]
+; CHECK-NEXT:    [[L_1_LCSSA3:%.*]] = phi ptr [ [[L_1]], [[INNER_LATCH]] ]
 ; CHECK-NEXT:    br label [[EXIT]]
 ; CHECK:       exit:
-; CHECK-NEXT:    [[L_15:%.*]] = phi ptr [ [[L_1_LCSSA4]], [[EXIT_LOOPEXIT1]] ], [ [[L_1_LCSSA]], [[EXIT_LOOPEXIT]] ]
-; CHECK-NEXT:    [[L_3:%.*]] = load i16, ptr [[L_15]], align 2
+; CHECK-NEXT:    [[L_14:%.*]] = phi ptr [ [[L_1_LCSSA3]], [[EXIT_LOOPEXIT1]] ], [ [[L_1_LCSSA]], [[EXIT_LOOPEXIT]] ]
+; CHECK-NEXT:    [[L_3:%.*]] = load i16, ptr [[L_14]], align 2
 ; CHECK-NEXT:    ret i16 [[L_3]]
 ;
 entry:
   br label %outer
 
 outer:
   %l.1 = load i16*, ptr %arg, align 8
   %l.2 = load i16*, ptr %arg, align 8