[llvm] (Draft) [SCEV] forgetValue: support (extractvalue 0, (with-overflow-i… (PR #98015)

[via llvm-commits]

https://github.com/v01dXYZ created https://github.com/llvm/llvm-project/pull/98015

Starter for fixing #97586 

>From 45e55186f48292ac7e9d82687ef49d724b679790 Mon Sep 17 00:00:00 2001
From: v01dxyz <v01dxyz at v01d.xyz>
Date: Mon, 8 Jul 2024 12:19:52 +0200
Subject: [PATCH] (Draft) [SCEV] forgetValue: support (extractvalue 0,
 (with-overflow-inst op0, op1))

Without that, forgetValue stops at (with-overflow-inst op0, op1) while
thanks to MatchBinaryOp, SCEV creation considers %extractvalue as if
it is (op op0, op1).

Because of that, it creates an unclearable SCEV value that could
possibly turn out of sync after a transform is applied.

The commit is in Draft as the fix is not satisfactory (code
duplication, do we push EVO to Visited ?). Meant for discussion and
for adding a test.
---
 llvm/lib/Analysis/ScalarEvolution.cpp         | 29 ++++++++++++++-
 .../Analysis/ScalarEvolutionTest.cpp          | 36 +++++++++++++++++++
 2 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index 430e1c6d8f8c69..8139c7d92548ac 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -8412,8 +8412,35 @@ void ScalarEvolution::visitAndClearUsers(
     SmallVectorImpl<const SCEV *> &ToForget) {
   while (!Worklist.empty()) {
     Instruction *I = Worklist.pop_back_val();
-    if (!isSCEVable(I->getType()))
+    if (!isSCEVable(I->getType())) {
+      // detect if a user is matching the pattern
+      // extractvalue 0, (with-overflow-inst op1, op2))
+      auto *WO = dyn_cast<WithOverflowInst>(I);
+
+      if (!WO)
+        continue;
+
+      for (auto *WOUser : WO->users()) {
+        auto *EVO = dyn_cast<ExtractValueInst>(WOUser);
+
+        if (!EVO)
+          continue;
+
+        if (EVO->getNumIndices() != 1 || EVO->getIndices()[0] != 0)
+          continue;
+
+        ValueExprMapType::iterator It =
+            ValueExprMap.find_as(static_cast<Value *>(EVO));
+        if (It != ValueExprMap.end()) {
+          eraseValueFromMap(It->first);
+          ToForget.push_back(It->second);
+        }
+
+        PushDefUseChildren(EVO, Worklist, Visited);
+      }
+
       continue;
+    }
 
     ValueExprMapType::iterator It =
         ValueExprMap.find_as(static_cast<Value *>(I));
diff --git a/llvm/unittests/Analysis/ScalarEvolutionTest.cpp b/llvm/unittests/Analysis/ScalarEvolutionTest.cpp
index a7b3c5c404ab75..a6a5ffda3cb706 100644
--- a/llvm/unittests/Analysis/ScalarEvolutionTest.cpp
+++ b/llvm/unittests/Analysis/ScalarEvolutionTest.cpp
@@ -1589,4 +1589,40 @@ TEST_F(ScalarEvolutionsTest, ApplyLoopGuards) {
   });
 }
 
+TEST_F(ScalarEvolutionsTest, ForgetValueWithOverflowInst) {
+  LLVMContext C;
+  SMDiagnostic Err;
+  std::unique_ptr<Module> M = parseAssemblyString(
+      "declare { i32, i1 } @llvm.smul.with.overflow.i32(i32, i32) "
+      "define void @foo(i32 %i) { "
+      "entry: "
+      "  br label %loop.body "
+      "loop.body: "
+      "  %iv = phi i32 [ %iv.next, %loop.body ], [ 0, %entry ] "
+      "  %iv.next = add nsw i32 %iv, 1 "
+      "  %call = call {i32, i1} @llvm.smul.with.overflow.i32(i32 %iv, i32 -2) "
+      "  %extractvalue = extractvalue {i32, i1} %call, 0 "
+      "  %cmp = icmp eq i32 %iv.next, 16 "
+      "  br i1 %cmp, label %exit, label %loop.body "
+      "exit: "
+      "  ret void "
+      "} ",
+      Err, C);
+
+  ASSERT_TRUE(M && "Could not parse module?");
+  ASSERT_TRUE(!verifyModule(*M) && "Must have been well formed!");
+
+  runWithSE(*M, "foo", [](Function &F, LoopInfo &LI, ScalarEvolution &SE) {
+    auto *ExtractValue = getInstructionByName(F, "extractvalue");
+    auto *IV = getInstructionByName(F, "iv");
+
+    auto *ExtractValueScev = SE.getSCEV(ExtractValue);
+    EXPECT_NE(ExtractValueScev, nullptr);
+
+    SE.forgetValue(IV);
+    auto *ExtractValueScevForgotten = SE.getExistingSCEV(ExtractValue);
+    EXPECT_EQ(ExtractValueScevForgotten, nullptr);
+  });
+}
+
 }  // end namespace llvm