[SCEVExpander] Expand umin_seq using freeze

[nikic]

Summary

%x umin_seq %y is currently expanded to %x == 0 ? 0 : umin(%x, %y). This patch changes the expansion to umin(%x, freeze %y) instead (https://alive2.llvm.org/ce/z/wujUhp).

The motivating for this change are the test cases affected by D124910, where the freeze expansion ultimately produces better optimization results. This is largely because (%x umin_seq %y) == %x is a common expansion pattern, which reliably optimizes in freeze representation, but only sometimes with the zero comparison (in particular, if %x == 0 can fold to something else, we generally won't be able to cover reasonable code from this.)

[reames]

LGTM.

I don't have any strong opinion as to which form is a better canonical one, and am thus open to trying the alternate and seeing how it works out.

[lebedev.ri]

Hello.

(This is part of my ~6 months backlog
in which i was not really around,
so i'm commenting on these things
as i stumble into them.)

While it is good that this helps optimizations,
there is one thing that is not mentioned at all here: round-tripping.
SCEV does not recognize this new canonical expansion back as umin_seq,
which means subsequent SCEV usages will have less insight into the IR.

I think, this is salvageable, but given that we model freeze as a SCEVUnkown,
the solution will be controversial.

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

private:
   Value *FindValueInExprValueMap(const SCEV *S, const Instruction *InsertPt);
 
   Value *expand(const SCEV *S);
 
   /// Determine the most "relevant" loop for the given SCEV.
   const Loop *getRelevantLoop(const SCEV *);
 
   Value *expandMinMaxExpr(const SCEVNAryExpr *S, Intrinsic::ID IntrinID,
-                          Twine Name);
+                          Twine Name, bool IsSequential = false);
 
   Value *visitConstant(const SCEVConstant *S) { return S->getValue(); }
 
   Value *visitPtrToIntExpr(const SCEVPtrToIntExpr *S);
 
   Value *visitTruncateExpr(const SCEVTruncateExpr *S);
 
   Value *visitZeroExtendExpr(const SCEVZeroExtendExpr *S);

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

Value *SCEVExpander::visitSignExtendExpr(const SCEVSignExtendExpr *S) {
   Type *Ty = SE.getEffectiveSCEVType(S->getType());
   Value *V = expandCodeForImpl(
       S->getOperand(), SE.getEffectiveSCEVType(S->getOperand()->getType()),
       false);
   return Builder.CreateSExt(V, Ty);
 }
 
 Value *SCEVExpander::expandMinMaxExpr(const SCEVNAryExpr *S,
-                                      Intrinsic::ID IntrinID, Twine Name) {
+                                      Intrinsic::ID IntrinID, Twine Name,
+                                      bool IsSequential) {
   Value *LHS = expand(S->getOperand(S->getNumOperands() - 1));
   Type *Ty = LHS->getType();
+  if (IsSequential)
+    LHS = Builder.CreateFreeze(LHS);
   for (int i = S->getNumOperands() - 2; i >= 0; --i) {
     Value *RHS = expandCodeForImpl(S->getOperand(i), Ty, false);
+    if (IsSequential && i != 0)
+      RHS = Builder.CreateFreeze(RHS);
     Value *Sel;
     if (Ty->isIntegerTy())
       Sel = Builder.CreateIntrinsic(IntrinID, {Ty}, {LHS, RHS},
                                     /*FMFSource=*/nullptr, Name);
     else {
       Value *ICmp =
           Builder.CreateICmp(MinMaxIntrinsic::getPredicate(IntrinID), LHS, RHS);
       Sel = Builder.CreateSelect(ICmp, LHS, RHS, Name);
Value *SCEVExpander::visitSMinExpr(const SCEVSMinExpr *S) {
   return expandMinMaxExpr(S, Intrinsic::smin, "smin");
 }
 
 Value *SCEVExpander::visitUMinExpr(const SCEVUMinExpr *S) {
   return expandMinMaxExpr(S, Intrinsic::umin, "umin");
 }
 
 Value *SCEVExpander::visitSequentialUMinExpr(const SCEVSequentialUMinExpr *S) {
-  SmallVector<Value *> Ops;
-  for (const SCEV *Op : S->operands())
-    Ops.emplace_back(expand(Op));
-
-  Value *SaturationPoint =
-      MinMaxIntrinsic::getSaturationPoint(Intrinsic::umin, S->getType());
-
-  SmallVector<Value *> OpIsZero;
-  for (Value *Op : ArrayRef<Value *>(Ops).drop_back())
-    OpIsZero.emplace_back(Builder.CreateICmpEQ(Op, SaturationPoint));
-
-  Value *AnyOpIsZero = Builder.CreateLogicalOr(OpIsZero);
-
-  Value *NaiveUMin = expandMinMaxExpr(S, Intrinsic::umin, "umin");
-  return Builder.CreateSelect(AnyOpIsZero, SaturationPoint, NaiveUMin);
+  return expandMinMaxExpr(S, Intrinsic::umin, "umin", /*IsSequential*/true);
 }
 
 Value *SCEVExpander::expandCodeForImpl(const SCEV *SH, Type *Ty,
                                        Instruction *IP, bool Root) {
   setInsertPoint(IP);
   Value *V = expandCodeForImpl(SH, Ty, Root);
   return V;
 }

llvm/test/Transforms/IndVarSimplify/exit-count-select.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -indvars -S | FileCheck %s
 
 define i32 @logical_and_2ops(i32 %n, i32 %m) {
 ; CHECK-LABEL: @logical_and_2ops(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[UMIN:%.*]] = call i32 @llvm.umin.i32(i32 [[M:%.*]], i32 [[N:%.*]])
+; CHECK-NEXT:    [[TMP0:%.*]] = freeze i32 [[M:%.*]]
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    br i1 false, label [[LOOP]], label [[EXIT:%.*]]
 ; CHECK:       exit:
-; CHECK-NEXT:    [[TMP0:%.*]] = icmp eq i32 [[N]], 0
-; CHECK-NEXT:    [[TMP1:%.*]] = select i1 [[TMP0]], i32 0, i32 [[UMIN]]
-; CHECK-NEXT:    ret i32 [[TMP1]]
+; CHECK-NEXT:    [[UMIN:%.*]] = call i32 @llvm.umin.i32(i32 [[TMP0]], i32 [[N:%.*]])
+; CHECK-NEXT:    ret i32 [[UMIN]]
 ;
 entry:
   br label %loop
 loop:
   %i = phi i32 [0, %entry], [%i.next, %loop]
   %i.next = add i32 %i, 1
   %cond_p0 = icmp ult i32 %i, %n
   %cond_p1 = icmp ult i32 %i, %m
   %cond = select i1 %cond_p0, i1 %cond_p1, i1 false
   br i1 %cond, label %loop, label %exit
 exit:
   ret i32 %i
 }
 
 define i32 @logical_or_2ops(i32 %n, i32 %m) {
 ; CHECK-LABEL: @logical_or_2ops(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[UMIN:%.*]] = call i32 @llvm.umin.i32(i32 [[M:%.*]], i32 [[N:%.*]])
+; CHECK-NEXT:    [[TMP0:%.*]] = freeze i32 [[M:%.*]]
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    br i1 true, label [[EXIT:%.*]], label [[LOOP]]
 ; CHECK:       exit:
-; CHECK-NEXT:    [[TMP0:%.*]] = icmp eq i32 [[N]], 0
-; CHECK-NEXT:    [[TMP1:%.*]] = select i1 [[TMP0]], i32 0, i32 [[UMIN]]
-; CHECK-NEXT:    ret i32 [[TMP1]]
+; CHECK-NEXT:    [[UMIN:%.*]] = call i32 @llvm.umin.i32(i32 [[TMP0]], i32 [[N:%.*]])
+; CHECK-NEXT:    ret i32 [[UMIN]]
 ;
 entry:
   br label %loop
 loop:
   %i = phi i32 [0, %entry], [%i.next, %loop]
   %i.next = add i32 %i, 1
   %cond_p0 = icmp uge i32 %i, %n
   %cond_p1 = icmp uge i32 %i, %m
   %cond = select i1 %cond_p0, i1 true, i1 %cond_p1
   br i1 %cond, label %exit, label %loop
 exit:
   ret i32 %i
 }
 
 define i32 @logical_and_3ops(i32 %n, i32 %m, i32 %k) {
 ; CHECK-LABEL: @logical_and_3ops(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[TMP0:%.*]] = icmp eq i32 [[M:%.*]], 0
-; CHECK-NEXT:    [[UMIN:%.*]] = call i32 @llvm.umin.i32(i32 [[K:%.*]], i32 [[M]])
-; CHECK-NEXT:    [[UMIN1:%.*]] = call i32 @llvm.umin.i32(i32 [[UMIN]], i32 [[N:%.*]])
+; CHECK-NEXT:    [[TMP0:%.*]] = freeze i32 [[K:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = freeze i32 [[M:%.*]]
+; CHECK-NEXT:    [[UMIN:%.*]] = call i32 @llvm.umin.i32(i32 [[TMP0]], i32 [[TMP1]])
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    br i1 false, label [[LOOP]], label [[EXIT:%.*]]
 ; CHECK:       exit:
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i32 [[N]], 0
-; CHECK-NEXT:    [[TMP2:%.*]] = select i1 [[TMP1]], i1 true, i1 [[TMP0]]
-; CHECK-NEXT:    [[TMP3:%.*]] = select i1 [[TMP2]], i32 0, i32 [[UMIN1]]
-; CHECK-NEXT:    ret i32 [[TMP3]]
+; CHECK-NEXT:    [[UMIN1:%.*]] = call i32 @llvm.umin.i32(i32 [[UMIN]], i32 [[N:%.*]])
+; CHECK-NEXT:    ret i32 [[UMIN1]]
 ;
 entry:
   br label %loop
 loop:
   %i = phi i32 [0, %entry], [%i.next, %loop]
   %i.next = add i32 %i, 1
   %cond_p0 = icmp ult i32 %i, %n
   %cond_p1 = icmp ult i32 %i, %m
   %cond_p2 = icmp ult i32 %i, %k
   %cond_p3 = select i1 %cond_p0, i1 %cond_p1, i1 false
   %cond = select i1 %cond_p3, i1 %cond_p2, i1 false
   br i1 %cond, label %loop, label %exit
 exit:
   ret i32 %i
 }
 
 define i32 @logical_or_3ops(i32 %n, i32 %m, i32 %k) {
 ; CHECK-LABEL: @logical_or_3ops(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[TMP0:%.*]] = icmp eq i32 [[M:%.*]], 0
-; CHECK-NEXT:    [[UMIN:%.*]] = call i32 @llvm.umin.i32(i32 [[K:%.*]], i32 [[M]])
-; CHECK-NEXT:    [[UMIN1:%.*]] = call i32 @llvm.umin.i32(i32 [[UMIN]], i32 [[N:%.*]])
+; CHECK-NEXT:    [[TMP0:%.*]] = freeze i32 [[K:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = freeze i32 [[M:%.*]]
+; CHECK-NEXT:    [[UMIN:%.*]] = call i32 @llvm.umin.i32(i32 [[TMP0]], i32 [[TMP1]])
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    br i1 true, label [[EXIT:%.*]], label [[LOOP]]
 ; CHECK:       exit:
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i32 [[N]], 0
-; CHECK-NEXT:    [[TMP2:%.*]] = select i1 [[TMP1]], i1 true, i1 [[TMP0]]
-; CHECK-NEXT:    [[TMP3:%.*]] = select i1 [[TMP2]], i32 0, i32 [[UMIN1]]
-; CHECK-NEXT:    ret i32 [[TMP3]]
+; CHECK-NEXT:    [[UMIN1:%.*]] = call i32 @llvm.umin.i32(i32 [[UMIN]], i32 [[N:%.*]])
+; CHECK-NEXT:    ret i32 [[UMIN1]]
 ;
 entry:
   br label %loop
 loop:
   %i = phi i32 [0, %entry], [%i.next, %loop]
   %i.next = add i32 %i, 1
   %cond_p0 = icmp uge i32 %i, %n
   %cond_p1 = icmp uge i32 %i, %m
   %cond_p2 = icmp uge i32 %i, %k
   %cond_p3 = select i1 %cond_p0, i1 true, i1 %cond_p1
   %cond = select i1 %cond_p3, i1 true, i1 %cond_p2
   br i1 %cond, label %exit, label %loop
 exit:
   ret i32 %i
 }