[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 c2a9a974ca85e4ac4509e368d4b9acae7e67bf71 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 (NFC)

---
 .../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 70091dc943ade280d75cea1e5ea5e93d9a8f934a 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 that leverage
LICM foldings (e.g. IndVars), may reason on invalid expressions; 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..5eccf7b4adb65 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