[LoopFlatten] Forget all block and loop dispositions after flatten

[StephenFan]

Summary

Method forgetLoop only forgets expression of phi or its users. SCEV
expressions except the above mentioned may still has loop dispositions
that point to the destroyed loop, which might cause a crash.

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

[SjoerdMeijer]

This looks good to me, but I would like to get a second opinion from @fhahn or @nikic just to be sure on the use of these scev helpers.

[StephenFan]

In D137651#3915315, @SjoerdMeijer wrote:

This looks good to me, but I would like to get a second opinion from @fhahn or @nikic just to be sure on the use of these scev helpers.

Thanks for taking a look!

[StephenFan]

Replace forgetValue as fogetBlockAndLoopDispositions.

[nikic]

Are you sure we need to forget the incoming value rather than the phi? I haven't looked at the code in detail, but I'd generally expect that removing an incoming value doesn't change dispositions of that value, but may change dispositions of the using phi.

llvm/test/Transforms/LoopFlatten/pr58865.ll

2

I'd suggest adding check lines as well.

[StephenFan]

In D137651#3915377, @nikic wrote:

Are you sure we need to forget the incoming value rather than the phi? I haven't looked at the code in detail, but I'd generally expect that removing an incoming value doesn't change dispositions of that value, but may change dispositions of the using phi.

If I understand correctly, the phi will be forgotten when calling forgetLoop method in the below. But I found that just forgetting the incoming value is still not sufficient or is wrong.

[StephenFan]

Replace forgetting incoming values as forget all of the block and loop dispositions.

[StephenFan]

Add checking lines.

[nikic]

LGTM

[fhahn]

LGTM, thanks!

[SjoerdMeijer]

This is (slightly) offtopic, but I am wondering what the lesson learned here is for me. I.e., I forgot forgetBlockAndLoopDispositions, but was wondering how I could/should have known.... I will check some SCEV documentation, that's one thing, but I am also wondering if SCEV could be made more robust. It was running in an assert because it thought it was in an invalid state, so could it call forgetBlockAndLoopDispositions instead of asserting?

[fhahn]

In D137651#3916997, @SjoerdMeijer wrote:

I will check some SCEV documentation, that's one thing, but I am also wondering if SCEV could be made more robust. It was running in an assert because it thought it was in an invalid state, so could it call forgetBlockAndLoopDispositions instead of asserting?

Unfortunately that's unlikely. We only catch this with verification, which basically computes most SCEV information again from scratch and compares it to the cached values, so it is super expensive.

[craig.topper]

Is this the same issue I previously reported in D109958? Did we lose that fix when that was reverted?

[SjoerdMeijer]

In D137651#3924013, @craig.topper wrote:

Is this the same issue I previously reported in D109958? Did we lose that fix when that was reverted?

Forgot to reply, but yeah, sorry, I think that was my bad. I think we accidentally lost that with a revert.

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

LLVM_DEBUG(dbgs() << "Replacing: "; V->dump(); dbgs() << "with:      ";
                OuterValue->dump());
     V->replaceAllUsesWith(OuterValue);
   }
 
   // Tell LoopInfo, SCEV and the pass manager that the inner loop has been
   // deleted, and any information that have about the outer loop invalidated.
   SE->forgetLoop(FI.OuterLoop);
   SE->forgetLoop(FI.InnerLoop);
+  SE->forgetBlockAndLoopDispositions();
   if (U)
     U->markLoopAsDeleted(*FI.InnerLoop, FI.InnerLoop->getName());
   LI->erase(FI.InnerLoop);
 
   // Increment statistic value.
   NumFlattened++;
 
   return true;

llvm/test/Transforms/LoopFlatten/pr58865.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -loop-flatten -verify-scev -S < %s | FileCheck %s
+;
+define void @sum_2d(ptr %p) {
+; CHECK-LABEL: @sum_2d(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[FLATTEN_TRIPCOUNT:%.*]] = mul i16 1, 1
+; CHECK-NEXT:    br label [[OUTER:%.*]]
+; CHECK:       outer:
+; CHECK-NEXT:    [[SUM_04:%.*]] = phi i16 [ 0, [[ENTRY:%.*]] ], [ [[DOTLCSSA:%.*]], [[OUTER_LATCH:%.*]] ]
+; CHECK-NEXT:    [[OUTER_IV:%.*]] = phi i16 [ 0, [[ENTRY]] ], [ [[OUTER_IV_NEXT:%.*]], [[OUTER_LATCH]] ]
+; CHECK-NEXT:    br label [[INNER:%.*]]
+; CHECK:       inner:
+; CHECK-NEXT:    [[SUM_12:%.*]] = phi i16 [ [[SUM_04]], [[OUTER]] ]
+; CHECK-NEXT:    [[INNER_IV:%.*]] = phi i16 [ 0, [[OUTER]] ]
+; CHECK-NEXT:    [[TMP0:%.*]] = load i16, ptr [[P:%.*]], align 1
+; CHECK-NEXT:    [[INNER_IV_NEXT:%.*]] = add nsw i16 [[INNER_IV]], 1
+; CHECK-NEXT:    [[CMP2:%.*]] = icmp slt i16 [[INNER_IV]], 0
+; CHECK-NEXT:    br label [[OUTER_LATCH]]
+; CHECK:       outer.latch:
+; CHECK-NEXT:    [[DOTLCSSA]] = phi i16 [ [[TMP0]], [[INNER]] ]
+; CHECK-NEXT:    [[OUTER_IV_NEXT]] = add nsw i16 [[OUTER_IV]], 1
+; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i16 [[OUTER_IV_NEXT]], [[FLATTEN_TRIPCOUNT]]
+; CHECK-NEXT:    br i1 [[CMP]], label [[OUTER]], label [[EXIT:%.*]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %outer
+
+outer:                                         ; preds = %outer.latch, %entry
+  %sum.04 = phi i16 [ 0, %entry ], [ %0, %outer.latch ]
+  %outer.iv = phi i16 [ 0, %entry ], [ %outer.iv.next, %outer.latch ]
+  br label %inner
+
+inner:                                        ; preds = %inner, %outer
+  %sum.12 = phi i16 [ %sum.04, %outer ], [ %0, %inner ]
+  %inner.iv = phi i16 [ 0, %outer ], [ %inner.iv.next, %inner ]
+  %0 = load i16, ptr %p, align 1
+  %inner.iv.next = add nsw i16 %inner.iv, 1
+  %cmp2 = icmp slt i16 %inner.iv, 0
+  br i1 %cmp2, label %inner, label %outer.latch
+
+outer.latch:                                         ; preds = %inner
+  %outer.iv.next = add nsw i16 %outer.iv, 1
+  %cmp = icmp slt i16 %outer.iv.next, 0
+  br i1 %cmp, label %outer, label %exit
+
+exit:                                         ; preds = %outer.latch
+  ret void
+}