[IndVarSimplify] Don't replace IV user with unsafe loop-invariant (PR45360)

[lebedev.ri]

Summary

As PR45360 reports,
with new cost-model we can sometimes end up being able to expand udiv/urem instructions.
And that exposes at least one instance of when we do that
regardless of whether or not it is safe to do.
In this particular case, it's SimplifyIndvar::replaceIVUserWithLoopInvariant().

It seems to me, we simply need to check with isSafeToExpandAt() first.

The test isn't great. I'm not sure how to make it only run -indvars.

Fixes PR45360.

[lebedev.ri]

Actually, i guess we only need isSafeToExpand(), not whole isSafeToExpandAt().

[mkazantsev]

Generally looks ok. Please use isSafeToExpandAt because the context can give us more information about divisor being non-zero.

llvm/lib/Transforms/Utils/SimplifyIndVar.cpp

660

This comment doesn't reflect what this method does, can you pls also fix it along, like, "Replace the UseInst with a loop invariant expression if it is safe"?

675

Please isSafeToExpandAt.

677

with non-speculable loop invariant maybe?

[lebedev.ri]

Thanks for taking a look!
Addressing review notes.

[mkazantsev]

LGTM, thanks!

[lebedev.ri]

In D82108#2108221, @mkazantsev wrote:

LGTM, thanks!

Thank you for the review!

llvm/lib/Transforms/Utils/SimplifyIndVar.cpp

 
 static Instruction *GetLoopInvariantInsertPosition(Loop *L, Instruction *Hint) {
   if (auto *BB = L->getLoopPreheader())
     return BB->getTerminator();
 
   return Hint;
 }
 
-/// Replace the UseInst with a constant if possible.
+/// Replace the UseInst with a loop invariant expression if it is safe.
 bool SimplifyIndvar::replaceIVUserWithLoopInvariant(Instruction *I) {
   if (!SE->isSCEVable(I->getType()))
     return false;
 
   // Get the symbolic expression for this instruction.
   const SCEV *S = SE->getSCEV(I);
 
   if (!SE->isLoopInvariant(S, L))
     return false;
 
   // Do not generate something ridiculous even if S is loop invariant.
   if (Rewriter.isHighCostExpansion(S, L, SCEVCheapExpansionBudget, TTI, I))
     return false;
 
   auto *IP = GetLoopInvariantInsertPosition(L, I);
+
+  if (!isSafeToExpandAt(S, IP, *SE)) {
+    LLVM_DEBUG(dbgs() << "INDVARS: Can not replace IV user: " << *I
+                      << " with non-speculable loop invariant: " << *S << '\n');
+    return false;
+  }
+
   auto *Invariant = Rewriter.expandCodeFor(S, I->getType(), IP);
 
   I->replaceAllUsesWith(Invariant);
   LLVM_DEBUG(dbgs() << "INDVARS: Replace IV user: " << *I
                     << " with loop invariant: " << *S << '\n');
   ++NumFoldedUser;
   Changed = true;
   DeadInsts.emplace_back(I);

llvm/test/Transforms/IndVarSimplify/X86/pr45360.ll

 ; CHECK-NEXT:  bb:
 ; CHECK-NEXT:    [[I6:%.*]] = load i32, i32* @a, align 4
 ; CHECK-NEXT:    [[I24:%.*]] = load i32, i32* @b, align 4
 ; CHECK-NEXT:    [[D_PROMOTED8:%.*]] = load i32, i32* @d, align 4
 ; CHECK-NEXT:    br label [[BB19:%.*]]
 ; CHECK:       bb19:
 ; CHECK-NEXT:    [[I8_LCSSA9:%.*]] = phi i32 [ [[D_PROMOTED8]], [[BB:%.*]] ], [ [[I8:%.*]], [[BB27:%.*]] ]
 ; CHECK-NEXT:    [[I8]] = and i32 [[I8_LCSSA9]], [[I6]]
-; CHECK-NEXT:    [[TMP0:%.*]] = urem i32 [[I24]], [[I8]]
 ; CHECK-NEXT:    [[I21:%.*]] = icmp eq i32 [[I8]], 0
 ; CHECK-NEXT:    br i1 [[I21]], label [[BB27_THREAD:%.*]], label [[BB27]]
 ; CHECK:       bb27.thread:
 ; CHECK-NEXT:    store i32 -1, i32* @f, align 4
 ; CHECK-NEXT:    store i32 0, i32* @d, align 4
 ; CHECK-NEXT:    store i32 0, i32* @c, align 4
 ; CHECK-NEXT:    br label [[BB32:%.*]]
 ; CHECK:       bb27:
-; CHECK-NEXT:    store i32 [[TMP0]], i32* @e, align 4
-; CHECK-NEXT:    [[I30:%.*]] = icmp eq i32 [[TMP0]], 0
+; CHECK-NEXT:    [[I26:%.*]] = urem i32 [[I24]], [[I8]]
+; CHECK-NEXT:    store i32 [[I26]], i32* @e, align 4
+; CHECK-NEXT:    [[I30:%.*]] = icmp eq i32 [[I26]], 0
 ; CHECK-NEXT:    br i1 [[I30]], label [[BB32_LOOPEXIT:%.*]], label [[BB19]]
 ; CHECK:       bb32.loopexit:
 ; CHECK-NEXT:    store i32 -1, i32* @f, align 4
 ; CHECK-NEXT:    store i32 [[I8]], i32* @d, align 4
 ; CHECK-NEXT:    br label [[BB32]]
 ; CHECK:       bb32:
 ; CHECK-NEXT:    [[C_SINK:%.*]] = phi i32* [ @c, [[BB32_LOOPEXIT]] ], [ @e, [[BB27_THREAD]] ]
 ; CHECK-NEXT:    store i32 0, i32* [[C_SINK]], align 4