[LoopReroll] Fix rerolling loop with use outside the loop

[kawashima-fj]

Summary

Fixes PR41696

The loop-reroll pass generates an invalid IR (or its assertion
fails in debug build) if values of the base instruction and
other root instructions (terms used in the loop-reroll pass)
are used outside the loop block. See IRs written in PR41696
as examples.

The current implementation of the loop-reroll pass can reroll
only loops that don't have values that are used outside the
loop, except reduced values (the last values of reduction chains).
This is described in the comment of the LoopReroll::reroll
function.
https://github.com/llvm/llvm-project/blob/llvmorg-10.0.0/llvm/lib/Transforms/Scalar/LoopRerollPass.cpp#L1600

This is checked in the LoopReroll::DAGRootTracker::validate
function.
https://github.com/llvm/llvm-project/blob/llvmorg-10.0.0/llvm/lib/Transforms/Scalar/LoopRerollPass.cpp#L1393

However, the base instruction and other root instructions skip
this check in the validation loop.
https://github.com/llvm/llvm-project/blob/llvmorg-10.0.0/llvm/lib/Transforms/Scalar/LoopRerollPass.cpp#L1229

Moving the check in front of the skip is the logically simplest
fix. However, inserting the check in an earlier stage is better
in terms of compilation time of unrerollable loops. This fix
inserts the check for the base instruction into the function
to validate possible base/root instructions. Check for other
root instructions is unnecessary because they don't match any
base instructions if they have uses outside the loop.

[efriedma]

I think there are four loop-reroll issues currently open in Bugzilla:

https://bugs.llvm.org/show_bug.cgi?id=34029
https://bugs.llvm.org/show_bug.cgi?id=42267
https://bugs.llvm.org/show_bug.cgi?id=34760
https://bugs.llvm.org/show_bug.cgi?id=41696

Could you check which ones this patch fixes?

llvm/lib/Transforms/Scalar/LoopRerollPass.cpp

882

Probably worth putting the "Check for other root instructions is unnecessary because they don't match any base instructions if they have uses outside the loop." explanation into a comment here.

[kawashima-fj]

Comment is added.

[kawashima-fj]

In D79549#2025775, @efriedma wrote:

I think there are four loop-reroll issues currently open in Bugzilla:

https://bugs.llvm.org/show_bug.cgi?id=34029
https://bugs.llvm.org/show_bug.cgi?id=42267
https://bugs.llvm.org/show_bug.cgi?id=34760
https://bugs.llvm.org/show_bug.cgi?id=41696

Could you check which ones this patch fixes?

I checked the following PRs. Unfortunately, only 41696, which I originally wanted to fix, is fixed.

https://bugs.llvm.org/show_bug.cgi?id=33085
https://bugs.llvm.org/show_bug.cgi?id=34029
https://bugs.llvm.org/show_bug.cgi?id=34639
https://bugs.llvm.org/show_bug.cgi?id=34760
https://bugs.llvm.org/show_bug.cgi?id=41696 (My target PR)
https://bugs.llvm.org/show_bug.cgi?id=42043
https://bugs.llvm.org/show_bug.cgi?id=42267
https://bugs.llvm.org/show_bug.cgi?id=42275
https://bugs.llvm.org/show_bug.cgi?id=42444

[efriedma]

Could you take another look at the testcase for 34760 in particular? That seems similar enough that it should get fixed at the same time.

[kawashima-fj]

In D79549#2027214, @efriedma wrote:

Could you take another look at the testcase for 34760 in particular? That seems similar enough that it should get fixed at the same time.

My change does not fix PR34760.
I tried the original reporter's precedure. LLVM bofore this change and after this change both give the same incorrect value (8).
I also tried opt -S -loop-reroll for the bitcode attached in the Comment 2. LLVM bofore this change and after this change both give same IR.
I read the bitcode attached in the Comment 2. It is not the pattern which I fixed.

In a quick glance, the problem of PR34760 is in the root finding process. I'll put a comment to PR34760 later.

[efriedma]

LGTM

llvm/lib/Transforms/Scalar/LoopRerollPass.cpp

for (User *V : I->users()) {
     // The recursive call makes a copy of SubsumedInsts.
     findRootsRecursive(I, SubsumedInsts);
   }
 }
 
 bool LoopReroll::DAGRootTracker::validateRootSet(DAGRootSet &DRS) {
   if (DRS.Roots.empty())
     return false;
+  if (hasUsesOutsideLoop(DRS.BaseInst, L))
+    return false;
 
   // Consider a DAGRootSet with N-1 roots (so N different values including
   //   BaseInst).
   // Define d = Roots[0] - BaseInst, which should be the same as
   //   Roots[I] - Roots[I-1] for all I in [1..N).
   // Define D = BaseInst@J - BaseInst@J-1, where "@J" means the value at the
   //   loop iteration J.
   //

llvm/test/Transforms/LoopReroll/external_use.ll

+; RUN: opt < %s -loop-reroll -S | FileCheck %s
+
+; Check whether rerolling is rejected if values of the base and root
+; instruction are used outside the loop block.
+
+; Only the base/root instructions except a loop increment instruction
+define void @test1() {
+entry:
+  br label %loop1
+
+loop1:
+;CHECK-LABEL: loop1:
+;CHECK-NEXT:   %indvar = phi i64 [ 0, %entry ], [ %indvar.next, %loop1 ]
+;CHECK-NEXT:   %indvar.1 = add nsw i64 %indvar, 1
+
+  %indvar = phi i64 [ 0, %entry ], [ %indvar.next, %loop1 ]
+  %indvar.1 = add nsw i64 %indvar, 1
+  %indvar.next = add nsw i64 %indvar, 2
+  %cmp = icmp slt i64 %indvar.next, 200
+  br i1 %cmp, label %loop1, label %exit
+
+exit:
+  %var1 = phi i64 [ %indvar.1, %loop1 ]
+  %var2 = phi i64 [ %indvar, %loop1 ]
+  ret void
+}
+
+; Both the base/root instructions and reduction instructions
+define void @test2() {
+entry:
+  br label %loop2
+
+loop2:
+;CHECK-LABEL: loop2:
+;CHECK-NEXT:   %indvar = phi i32  [ 0, %entry ], [ %indvar.next, %loop2 ]
+;CHECK-NEXT:   %redvar = phi i32 [ 0, %entry ], [ %add.2, %loop2 ]
+;CHECK-NEXT:   %indvar.1 = add nuw nsw i32 %indvar, 1
+;CHECK-NEXT:   %indvar.2 = add nuw nsw i32 %indvar, 2
+
+  %indvar = phi i32 [ 0, %entry ], [ %indvar.next, %loop2 ]
+  %redvar = phi i32 [ 0, %entry ], [ %add.2, %loop2 ]
+  %indvar.1 = add nuw nsw i32 %indvar, 1
+  %indvar.2 = add nuw nsw i32 %indvar, 2
+  %mul.0 = mul nsw i32 %indvar, %indvar
+  %mul.1 = mul nsw i32 %indvar.1, %indvar.1
+  %mul.2 = mul nsw i32 %indvar.2, %indvar.2
+  %add.0 = add nsw i32 %redvar, %mul.0
+  %add.1 = add nsw i32 %add.0, %mul.1
+  %add.2 = add nsw i32 %add.1, %mul.2
+  %indvar.next = add nuw nsw i32 %indvar, 3
+  %cmp = icmp slt i32 %indvar.next, 300
+  br i1 %cmp, label %loop2, label %exit
+
+exit:
+  %a = phi i32 [ %indvar, %loop2 ]
+  %b = phi i32 [ %indvar.1, %loop2 ]
+  %c = phi i32 [ %indvar.2, %loop2 ]
+  %x = phi i32 [ %add.2, %loop2 ]
+  ret void
+}