From a8e9742bd44daa4632b07ccc37a6da8236781fa6 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo@fhahn.com>
Date: Tue, 18 Oct 2022 16:17:49 +0100
Subject: [PATCH] [IndVarSimplify] Clear block and loop dispositions after
 moving instr.

Moving an instruction can invalidate the cached block dispositions of
the corresponding SCEV. Invalidate the cached dispositions.

Also fixes a copy-paste error in forgetBlockAndLoopDispositions where
the start expression S was removed from BlockDispositions in the loop
but not the current values. This was also exposed by the new test case.

Fixes #58439.
---
 llvm/lib/Analysis/ScalarEvolution.cpp              |  5 ++--
 llvm/lib/Transforms/Scalar/IndVarSimplify.cpp      |  1 +
 .../{pr54434.ll => scev-invalidation.ll}           | 33 ++++++++++++++++++++--
 3 files changed, 34 insertions(+), 5 deletions(-)
 rename llvm/test/Transforms/IndVarSimplify/{pr54434.ll => scev-invalidation.ll} (62%)

diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index d8ba9c6..985087c 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -8408,9 +8408,10 @@ void ScalarEvolution::forgetBlockAndLoopDispositions(Value *V) {
   SmallPtrSet<const SCEV *, 8> Seen = {S};
   while (!Worklist.empty()) {
     const SCEV *Curr = Worklist.pop_back_val();
-    if (!LoopDispositions.erase(Curr) && !BlockDispositions.erase(S))
+    bool LoopDispoRemoved = LoopDispositions.erase(Curr);
+    bool BlockDispoRemoved = BlockDispositions.erase(Curr);
+    if (!LoopDispoRemoved && !BlockDispoRemoved)
       continue;
-
     auto Users = SCEVUsers.find(Curr);
     if (Users != SCEVUsers.end())
       for (const auto *User : Users->second)
diff --git a/llvm/lib/Transforms/Scalar/IndVarSimplify.cpp b/llvm/lib/Transforms/Scalar/IndVarSimplify.cpp
index 42a3ae2..9efb40f 100644
--- a/llvm/lib/Transforms/Scalar/IndVarSimplify.cpp
+++ b/llvm/lib/Transforms/Scalar/IndVarSimplify.cpp
@@ -1281,6 +1281,7 @@ bool IndVarSimplify::sinkUnusedInvariants(Loop *L) {
 
     MadeAnyChanges = true;
     ToMove->moveBefore(*ExitBlock, InsertPt);
+    SE->forgetBlockAndLoopDispositions(ToMove);
     if (Done) break;
     InsertPt = ToMove->getIterator();
   }
diff --git a/llvm/test/Transforms/IndVarSimplify/pr54434.ll b/llvm/test/Transforms/IndVarSimplify/scev-invalidation.ll
similarity index 62%
rename from llvm/test/Transforms/IndVarSimplify/pr54434.ll
rename to llvm/test/Transforms/IndVarSimplify/scev-invalidation.ll
index 7f25c6d..0dac4f3 100644
--- a/llvm/test/Transforms/IndVarSimplify/pr54434.ll
+++ b/llvm/test/Transforms/IndVarSimplify/scev-invalidation.ll
@@ -1,8 +1,8 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
-; RUN: opt -S -indvars -verify-scev < %s | FileCheck %s
+; RUN: opt -S -passes=indvars -verify-scev < %s | FileCheck %s
 
-define void @test() {
-; CHECK-LABEL: @test(
+define void @test_pr54434() {
+; CHECK-LABEL: @test_pr54434(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[FOR_COND:%.*]]
 ; CHECK:       for.cond:
@@ -43,3 +43,30 @@ for.end106:                                       ; preds = %for.cond
   ret void
 }
 
+define i32 @test_pr58439(i32 %a) {
+; CHECK-LABEL: @test_pr58439(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    br i1 false, label [[LOOP]], label [[EXIT:%.*]]
+; CHECK:       exit:
+; CHECK-NEXT:    [[C_EXT_LCSSA:%.*]] = phi i32 [ 0, [[LOOP]] ]
+; CHECK-NEXT:    [[OR:%.*]] = or i32 [[A:%.*]], 1
+; CHECK-NEXT:    [[RES:%.*]] = add i32 [[C_EXT_LCSSA]], [[OR]]
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+entry:
+  %or = or i32 %a, 1
+  br label %loop
+
+loop:
+  %iv = phi i32 [ 1, %entry ], [ %iv.next, %loop ]
+  %iv.next = add i32 %iv, 1
+  %c.1 = icmp eq i32 %iv.next, %or
+  %c.ext = zext i1 %c.1 to i32
+  br i1 false, label %loop, label %exit
+
+exit:
+  %res = add i32 %c.ext, %or
+  ret i32 %res
+}
-- 
2.7.4