[IndVarSimplify] Don't assert that terminator is not SCEVable (PR55925)

[nikic]

Summary

The IV widening code currently asserts that terminators aren't SCEVable -- however, this is not the case for invokes with a returned attribute.

As far as I can tell, this assertions is not necessary -- even if we have a critical edge (the second test case), the trunc gets inserted in a legal position. Though I would be happy to instead bail out for the terminator case, as it's something of an edge case, and I don't think the transform is profitable for it anyway.

Fixes https://github.com/llvm/llvm-project/issues/55925.

[fhahn]

LGTM, thanks! I think the utilities should be able to handle critical edges here now.

[ellis]

In D127288#3567502, @fhahn wrote:

LGTM, thanks! I think the utilities should be able to handle critical edges here now.

Hi @fhahn. Do you know if there are any recent commits that add the ability to handle this? I'm wondering if this work was recent, or if this commit can be used with an earlier version of LLVM. Thanks!

llvm/lib/Transforms/Utils/SimplifyIndVar.cpp

if (widenWithVariantUse(DU))
       return nullptr;
 
     // This user does not evaluate to a recurrence after widening, so don't
     // follow it. Instead insert a Trunc to kill off the original use,
     // eventually isolating the original narrow IV so it can be removed.
     truncateIVUse(DU, DT, LI);
     return nullptr;
   }
-  // Assume block terminators cannot evaluate to a recurrence. We can't to
-  // insert a Trunc after a terminator if there happens to be a critical edge.
-  assert(DU.NarrowUse != DU.NarrowUse->getParent()->getTerminator() &&
-         "SCEV is not expected to evaluate a block terminator");
 
   // Reuse the IV increment that SCEVExpander created as long as it dominates
   // NarrowUse.
   Instruction *WideUse = nullptr;
   if (WideAddRec.first == WideIncExpr &&
       Rewriter.hoistIVInc(WideInc, DU.NarrowUse))
     WideUse = WideInc;
   else {

llvm/test/Transforms/IndVarSimplify/pr55925.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -indvars < %s | FileCheck %s
+
+; This tests the case where a terminator can be modeled by SCEV,
+; because it has a returned attribute.
+
+target datalayout = "e-m:o-i64:64-i128:128-n32:64-S128"
+
+declare i32 @foo(i32)
+
+define void @test(i8* %p) personality i8* undef {
+; CHECK-LABEL: @test(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[INDVARS_IV:%.*]] = phi i64 [ [[INDVARS_IV_NEXT:%.*]], [[LOOP_LATCH:%.*]] ], [ 0, [[ENTRY:%.*]] ]
+; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ 0, [[ENTRY]] ], [ [[IV_NEXT:%.*]], [[LOOP_LATCH]] ]
+; CHECK-NEXT:    [[RES:%.*]] = invoke i32 @foo(i32 returned [[IV]])
+; CHECK-NEXT:    to label [[LOOP_LATCH]] unwind label [[EXIT:%.*]]
+; CHECK:       loop.latch:
+; CHECK-NEXT:    [[INDVARS_IV_NEXT]] = add i64 [[INDVARS_IV]], 1
+; CHECK-NEXT:    [[IV_NEXT]] = add nuw i32 [[IV]], 1
+; CHECK-NEXT:    [[TMP0:%.*]] = trunc i64 [[INDVARS_IV]] to i32
+; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @foo(i32 [[TMP0]])
+; CHECK-NEXT:    br label [[LOOP]]
+; CHECK:       exit:
+; CHECK-NEXT:    [[LP:%.*]] = landingpad { i8*, i32 }
+; CHECK-NEXT:    cleanup
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %loop
+
+loop:
+  %iv = phi i32 [ 0, %entry ], [ %iv.next, %loop.latch ]
+  %res = invoke i32 @foo(i32 returned %iv)
+  to label %loop.latch unwind label %exit
+
+loop.latch:
+  %ext = zext i32 %iv to i64
+  %tmp5 = getelementptr inbounds i8, i8* %p, i64 %ext
+  %iv.next = add nuw i32 %iv, 1
+  call i32 @foo(i32 %res)
+  br label %loop
+
+exit:
+  %lp = landingpad { i8*, i32 }
+  cleanup
+  ret void
+}
+
+define void @test_critedge(i1 %c, i8* %p) personality i8* undef {
+; CHECK-LABEL: @test_critedge(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[INDVARS_IV:%.*]] = phi i64 [ [[INDVARS_IV_NEXT:%.*]], [[LOOP_LATCH:%.*]] ], [ 0, [[ENTRY:%.*]] ]
+; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ 0, [[ENTRY]] ], [ [[IV_NEXT:%.*]], [[LOOP_LATCH]] ]
+; CHECK-NEXT:    br i1 [[C:%.*]], label [[LOOP_INVOKE:%.*]], label [[LOOP_OTHER:%.*]]
+; CHECK:       loop.invoke:
+; CHECK-NEXT:    [[TMP0:%.*]] = trunc i64 [[INDVARS_IV]] to i32
+; CHECK-NEXT:    [[RES:%.*]] = invoke i32 @foo(i32 returned [[IV]])
+; CHECK-NEXT:    to label [[LOOP_LATCH]] unwind label [[EXIT:%.*]]
+; CHECK:       loop.other:
+; CHECK-NEXT:    br label [[LOOP_LATCH]]
+; CHECK:       loop.latch:
+; CHECK-NEXT:    [[PHI:%.*]] = phi i32 [ [[TMP0]], [[LOOP_INVOKE]] ], [ 0, [[LOOP_OTHER]] ]
+; CHECK-NEXT:    [[INDVARS_IV_NEXT]] = add i64 [[INDVARS_IV]], 1
+; CHECK-NEXT:    [[IV_NEXT]] = add nuw i32 [[IV]], 1
+; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @foo(i32 [[PHI]])
+; CHECK-NEXT:    br label [[LOOP]]
+; CHECK:       exit:
+; CHECK-NEXT:    [[LP:%.*]] = landingpad { i8*, i32 }
+; CHECK-NEXT:    cleanup
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %loop
+
+loop:
+  %iv = phi i32 [ 0, %entry ], [ %iv.next, %loop.latch ]
+  br i1 %c, label %loop.invoke, label %loop.other
+
+loop.invoke:
+  %res = invoke i32 @foo(i32 returned %iv)
+  to label %loop.latch unwind label %exit
+
+loop.other:
+  br label %loop.latch
+
+loop.latch:
+  %phi = phi i32 [ %res, %loop.invoke ], [ 0, %loop.other ]
+  %ext = zext i32 %iv to i64
+  %tmp5 = getelementptr inbounds i8, i8* %p, i64 %ext
+  %iv.next = add nuw i32 %iv, 1
+  call i32 @foo(i32 %phi)
+  br label %loop
+
+exit:
+  %lp = landingpad { i8*, i32 }
+  cleanup
+  ret void
+}