[llvm] [LICM] Invalidate cached SCEV results after reassociation (PR #92655)

[Antonio Frighetto via llvm-commits]

https://github.com/antoniofrighetto updated https://github.com/llvm/llvm-project/pull/92655

>From 466ab00f7affce84b4ce8e6c9d0d9480f4dfab08 Mon Sep 17 00:00:00 2001
From: Antonio Frighetto <me at antoniofrighetto.com>
Date: Sat, 18 May 2024 16:46:03 +0200
Subject: [PATCH 1/2] [LICM] Introduce test for PR92655

---
 .../LICM/update-scev-after-hoist.ll           | 24 +++++++++++++++++++
 1 file changed, 24 insertions(+)
 create mode 100644 llvm/test/Transforms/LICM/update-scev-after-hoist.ll

diff --git a/llvm/test/Transforms/LICM/update-scev-after-hoist.ll b/llvm/test/Transforms/LICM/update-scev-after-hoist.ll
new file mode 100644
index 0000000000000..f834a74b6f247
--- /dev/null
+++ b/llvm/test/Transforms/LICM/update-scev-after-hoist.ll
@@ -0,0 +1,24 @@
+; RUN: opt -S -passes='loop-unroll,loop-mssa(licm),print<scalar-evolution>' -unroll-count=4 -disable-output < %s 2>&1 | FileCheck %s --check-prefix=SCEV-EXPR
+
+define i16 @main() {
+; SCEV-EXPR:      Classifying expressions for: @main
+; SCEV-EXPR-NEXT:  %mul = phi i16 [ 1, %entry ], [ %mul.n.3, %loop ]
+; SCEV-EXPR-NEXT:  -->  %mul U: full-set S: [-32768,32753)		Exits: 4096		LoopDispositions: { %loop: Variant }
+; SCEV-EXPR-NEXT:  %div = phi i16 [ 32767, %entry ], [ %div.n.3, %loop ]
+; SCEV-EXPR-NEXT:  -->  %div U: [-2048,-32768) S: [-2048,-32768)		Exits: 7		LoopDispositions: { %loop: Variant }
+; SCEV-EXPR-NEXT:  %mul.n = mul i16 %mul, 8
+; SCEV-EXPR-NEXT:  -->  (2 * %mul) U: [0,-1) S: [-32768,32767)		Exits: 8192		LoopDispositions: { %loop: Variant }
+entry:
+  br label %loop
+
+loop:
+  %mul = phi i16 [ 1, %entry ], [ %mul.n, %loop ]
+  %div = phi i16 [ 32767, %entry ], [ %div.n, %loop ]
+  %mul.n = mul i16 %mul, 2
+  %div.n = sdiv i16 %div, 2
+  %cmp = icmp sgt i16 %div, 0
+  br i1 %cmp, label %loop, label %end
+
+end:
+  ret i16 %mul
+}

>From 115a3e59f25b527d5caefddd005c1a0b4e57a48d Mon Sep 17 00:00:00 2001
From: Antonio Frighetto <me at antoniofrighetto.com>
Date: Sat, 18 May 2024 16:52:17 +0200
Subject: [PATCH 2/2] [LICM] Invalidate cached SCEV results in
 `hoistMulAddAssociation`

While reassociating expressions, LICM is required to invalidate SCEV
results, as otherwise subsequent passes in the pipeline (e.g., IndVars)
may reason on invalid results, thus miscompiling. This is achieved by
rewriting the reassociable instruction from scratch.

Fixes: https://github.com/llvm/llvm-project/issues/91957.
---
 llvm/lib/Transforms/Scalar/LICM.cpp              | 16 ++++++++++++++--
 .../Transforms/LICM/update-scev-after-hoist.ll   |  6 +++---
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/LICM.cpp b/llvm/lib/Transforms/Scalar/LICM.cpp
index 6aa4188d1cc4d..48cfff24791bd 100644
--- a/llvm/lib/Transforms/Scalar/LICM.cpp
+++ b/llvm/lib/Transforms/Scalar/LICM.cpp
@@ -2751,7 +2751,7 @@ static bool hoistMulAddAssociation(Instruction &I, Loop &L,
   IRBuilder<> Builder(Preheader->getTerminator());
   for (auto *U : Changes) {
     assert(L.isLoopInvariant(U->get()));
-    Instruction *Ins = cast<Instruction>(U->getUser());
+    auto *Ins = cast<BinaryOperator>(U->getUser());
     Value *Mul;
     if (I.getType()->isIntOrIntVectorTy()) {
       Mul = Builder.CreateMul(U->get(), Factor, "factor.op.mul");
@@ -2759,8 +2759,20 @@ static bool hoistMulAddAssociation(Instruction &I, Loop &L,
       Ins->dropPoisonGeneratingFlags();
     } else
       Mul = Builder.CreateFMulFMF(U->get(), Factor, Ins, "factor.op.fmul");
-    U->set(Mul);
+
+    // Rewrite the reassociable instruction
+    unsigned OpIdx = U->getOperandNo();
+    auto *LHS = OpIdx == 0 ? Mul : Ins->getOperand(0);
+    auto *RHS = OpIdx == 1 ? Mul : Ins->getOperand(1);
+    auto *NewBO = BinaryOperator::Create(Ins->getOpcode(), LHS, RHS,
+                                         Ins->getName() + ".reass", Ins);
+    NewBO->copyIRFlags(Ins);
+    if (VariantOp == Ins)
+      VariantOp = NewBO;
+    Ins->replaceAllUsesWith(NewBO);
+    eraseInstruction(*Ins, SafetyInfo, MSSAU);
   }
+
   I.replaceAllUsesWith(VariantOp);
   eraseInstruction(I, SafetyInfo, MSSAU);
   return true;
diff --git a/llvm/test/Transforms/LICM/update-scev-after-hoist.ll b/llvm/test/Transforms/LICM/update-scev-after-hoist.ll
index f834a74b6f247..fc45b8fce1766 100644
--- a/llvm/test/Transforms/LICM/update-scev-after-hoist.ll
+++ b/llvm/test/Transforms/LICM/update-scev-after-hoist.ll
@@ -3,11 +3,11 @@
 define i16 @main() {
 ; SCEV-EXPR:      Classifying expressions for: @main
 ; SCEV-EXPR-NEXT:  %mul = phi i16 [ 1, %entry ], [ %mul.n.3, %loop ]
-; SCEV-EXPR-NEXT:  -->  %mul U: full-set S: [-32768,32753)		Exits: 4096		LoopDispositions: { %loop: Variant }
+; SCEV-EXPR-NEXT:  -->  %mul U: [0,-15) S: [-32768,32753)		Exits: 4096		LoopDispositions: { %loop: Variant }
 ; SCEV-EXPR-NEXT:  %div = phi i16 [ 32767, %entry ], [ %div.n.3, %loop ]
 ; SCEV-EXPR-NEXT:  -->  %div U: [-2048,-32768) S: [-2048,-32768)		Exits: 7		LoopDispositions: { %loop: Variant }
-; SCEV-EXPR-NEXT:  %mul.n = mul i16 %mul, 8
-; SCEV-EXPR-NEXT:  -->  (2 * %mul) U: [0,-1) S: [-32768,32767)		Exits: 8192		LoopDispositions: { %loop: Variant }
+; SCEV-EXPR-NEXT:  %mul.n.reass.reass = mul i16 %mul, 8
+; SCEV-EXPR-NEXT:  -->  (8 * %mul) U: [0,-7) S: [-32768,32761)		Exits: -32768		LoopDispositions: { %loop: Variant }
 entry:
   br label %loop