[InstCombine] limit icmp fold with sub if other sub user is a phi

[spatel]

Summary

This is a hacky fix for issue #54558. As discussed there, codegen regressed when we opened up this transform to allow extra uses ( 61580d0949fd3465 ), and it's not clear how to undo the transforms at the later stage of compilation.

As noted in the code comments, there's a set of remaining folds that are still limited to one-use, so we can try harder to refine and expand the limitations on these folds, but it's likely to be an up-and-down battle as we find and overcome similar regressions.

[kazu]

Sanjay, thank you so much for the fix!

[lebedev.ri]

It seems like something for a late undo pass to deal with.

[spatel]

In D122909#3422897, @lebedev.ri wrote:

It seems like something for a late undo pass to deal with.

Agree, but the comments in https://github.com/llvm/llvm-project/issues/54558 suggested there wasn't a quick fix, and there are already one-use limitations here, so that's why I proposed a fast fix for the regression. Do you object to this change?

[lebedev.ri]

In D122909#3424859, @spatel wrote:

In D122909#3422897, @lebedev.ri wrote:

It seems like something for a late undo pass to deal with.

Agree, but the comments in https://github.com/llvm/llvm-project/issues/54558 suggested there wasn't a quick fix, and there are already one-use limitations here, so that's why I proposed a fast fix for the regression. Do you object to this change?

No strong thoughts on this one. I just think it would be best to have a more general undo fold :)

[spatel]

In D122909#3424861, @lebedev.ri wrote:

In D122909#3424859, @spatel wrote:

In D122909#3422897, @lebedev.ri wrote:

It seems like something for a late undo pass to deal with.

Agree, but the comments in https://github.com/llvm/llvm-project/issues/54558 suggested there wasn't a quick fix, and there are already one-use limitations here, so that's why I proposed a fast fix for the regression. Do you object to this change?

No strong thoughts on this one. I just think it would be best to have a more general undo fold :)

Yes, I didn't look closely at the original example, but there was a comment about a cross-block limitation, so it might be a job for CGP.

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

Instruction *InstCombinerImpl::foldICmpSubConstant(ICmpInst &Cmp,
   bool HasNUW = Sub->hasNoUnsignedWrap();
   if (match(X, m_APInt(C2)) &&
       ((Cmp.isUnsigned() && HasNUW) || (Cmp.isSigned() && HasNSW)) &&
       !subWithOverflow(SubResult, *C2, C, Cmp.isSigned()))
     return new ICmpInst(SwappedPred, Y, ConstantInt::get(Ty, SubResult));
 
   // X - Y == 0 --> X == Y.
   // X - Y != 0 --> X != Y.
-  if (Cmp.isEquality() && C.isZero())
+  // TODO: We allow this with multiple uses as long as the other uses are not
+  //       in phis. The phi use check is guarding against a codegen regression
+  //       for a loop test. If the backend could undo this (and possibly
+  //       subsequent transforms), we would not need this hack.
+  if (Cmp.isEquality() && C.isZero() &&
+      none_of((Sub->users()), [](const User *U) { return isa<PHINode>(U); }))
     return new ICmpInst(Pred, X, Y);
 
   // The following transforms are only worth it if the only user of the subtract
   // is the icmp.
   // TODO: This is an artificial restriction for all of the transforms below
-  //       that only need a single replacement icmp.
+  //       that only need a single replacement icmp. Can these use the phi test
+  //       like the transform above here?
   if (!Sub->hasOneUse())
     return nullptr;
 
   if (Sub->hasNoSignedWrap()) {
     // (icmp sgt (sub nsw X, Y), -1) -> (icmp sge X, Y)
     if (Pred == ICmpInst::ICMP_SGT && C.isAllOnes())
       return new ICmpInst(ICmpInst::ICMP_SGE, X, Y);
 

llvm/test/Transforms/InstCombine/icmp-sub.ll

 ;
   %sub = sub i32 %a, %b
   store i32 %sub, i32* %p
   %cmp = icmp eq i32 %sub, 0
   %sel = select i1 %cmp, i32 %a, i32 %b
   ret i32 %sel
 }
 
-; TODO: Replacing the "SUB == 0" regresses codegen, and it may be hard to recover from that.
+; Replacing the "SUB == 0" regresses codegen, and it may be hard to recover from that.
 
 declare i32 @llvm.umin.i32(i32, i32)
 
 define void @PR54558_reduced(i32 %arg) {
 ; CHECK-LABEL: @PR54558_reduced(
 ; CHECK-NEXT:  bb_entry:
 ; CHECK-NEXT:    br label [[BB_LOOP:%.*]]
 ; CHECK:       bb_loop:
 ; CHECK-NEXT:    [[PHI_OUTER:%.*]] = phi i32 [ [[SUB:%.*]], [[BB_LOOP]] ], [ [[ARG:%.*]], [[BB_ENTRY:%.*]] ]
 ; CHECK-NEXT:    [[MIN:%.*]] = tail call i32 @llvm.umin.i32(i32 [[PHI_OUTER]], i32 43)
 ; CHECK-NEXT:    call void @use(i32 [[MIN]])
 ; CHECK-NEXT:    [[SUB]] = sub i32 [[PHI_OUTER]], [[MIN]]
-; CHECK-NEXT:    [[COND_OUTER:%.*]] = icmp ult i32 [[PHI_OUTER]], 44
+; CHECK-NEXT:    [[COND_OUTER:%.*]] = icmp eq i32 [[SUB]], 0
 ; CHECK-NEXT:    br i1 [[COND_OUTER]], label [[BB_EXIT:%.*]], label [[BB_LOOP]]
 ; CHECK:       bb_exit:
 ; CHECK-NEXT:    ret void
 ;
 bb_entry:
   br label %bb_loop
 
 bb_loop:
   %phi_outer = phi i32 [ %sub, %bb_loop ], [ %arg, %bb_entry ]
   %min = tail call i32 @llvm.umin.i32(i32 %phi_outer, i32 43)
   call void @use(i32 %min)
   %sub = sub i32 %phi_outer, %min
   %cond_outer = icmp eq i32 %sub, 0
   br i1 %cond_outer, label %bb_exit, label %bb_loop
 
 bb_exit:
   ret void
 }
 
+; TODO: It might be ok to replace the "SUB == 0" in this example if codegen can invert it.
+
 define void @PR54558_reduced_more(i32 %x, i32 %y) {
 ; CHECK-LABEL: @PR54558_reduced_more(
 ; CHECK-NEXT:  bb_entry:
 ; CHECK-NEXT:    br label [[BB_LOOP:%.*]]
 ; CHECK:       bb_loop:
 ; CHECK-NEXT:    [[PHI_OUTER:%.*]] = phi i32 [ [[SUB:%.*]], [[BB_LOOP]] ], [ [[X:%.*]], [[BB_ENTRY:%.*]] ]
 ; CHECK-NEXT:    [[SUB]] = sub i32 [[PHI_OUTER]], [[Y:%.*]]
-; CHECK-NEXT:    [[COND_OUTER:%.*]] = icmp eq i32 [[PHI_OUTER]], [[Y]]
+; CHECK-NEXT:    [[COND_OUTER:%.*]] = icmp eq i32 [[SUB]], 0
 ; CHECK-NEXT:    br i1 [[COND_OUTER]], label [[BB_EXIT:%.*]], label [[BB_LOOP]]
 ; CHECK:       bb_exit:
 ; CHECK-NEXT:    ret void
 ;
 bb_entry:
   br label %bb_loop
 
 bb_loop:
   %phi_outer = phi i32 [ %sub, %bb_loop ], [ %x, %bb_entry ]
   %sub = sub i32 %phi_outer, %y
   %cond_outer = icmp eq i32 %sub, 0
   br i1 %cond_outer, label %bb_exit, label %bb_loop
 
 bb_exit:
   ret void
 }