[llvm] 70091dc - [LICM] Invalidate cached SCEV results in `hoistMulAddAssociation`

[Antonio Frighetto via llvm-commits]

Author: Antonio Frighetto
Date: 2024-05-29T08:44:45+02:00
New Revision: 70091dc943ade280d75cea1e5ea5e93d9a8f934a

URL: https://github.com/llvm/llvm-project/commit/70091dc943ade280d75cea1e5ea5e93d9a8f934a
DIFF: https://github.com/llvm/llvm-project/commit/70091dc943ade280d75cea1e5ea5e93d9a8f934a.diff

LOG: [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.

Added: 
    

Modified: 
    llvm/lib/Transforms/Scalar/LICM.cpp
    llvm/test/Transforms/LICM/update-scev-after-hoist.ll

Removed: 
    


################################################################################
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