[llvm] 161bfa5 - [LoopFlattening] Check for extra uses on Mul

[David Green via llvm-commits]

Author: David Green
Date: 2023-01-06T15:32:38Z
New Revision: 161bfa5f53efe5cb5a7b10e2eb21023c6a33a18b

URL: https://github.com/llvm/llvm-project/commit/161bfa5f53efe5cb5a7b10e2eb21023c6a33a18b
DIFF: https://github.com/llvm/llvm-project/commit/161bfa5f53efe5cb5a7b10e2eb21023c6a33a18b.diff

LOG: [LoopFlattening] Check for extra uses on Mul

Similar to D138404, we were not guarding against extra uses of the Mul.
In most cases other checks would catch the issue due to unsupported
instructions in the outer loop, but certain non-canonical loop forms
could still get through.

Fixes #59339

Differential Revision: https://reviews.llvm.org/D141114

Added: 
    

Modified: 
    llvm/lib/Transforms/Scalar/LoopFlatten.cpp
    llvm/test/Transforms/LoopFlatten/pr59339.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Scalar/LoopFlatten.cpp b/llvm/lib/Transforms/Scalar/LoopFlatten.cpp
index 9b6750bb4ea7..d82624ef9143 100644
--- a/llvm/lib/Transforms/Scalar/LoopFlatten.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopFlatten.cpp
@@ -215,6 +215,15 @@ struct FlattenInfo {
     LLVM_DEBUG(dbgs() << "Matched multiplication: "; MatchedMul->dump());
     LLVM_DEBUG(dbgs() << "Matched iteration count: "; MatchedItCount->dump());
 
+    // The mul should not have any other uses. Widening may leave trivially dead
+    // uses, which can be ignored.
+    if (count_if(MatchedMul->users(), [](User *U) {
+          return !isInstructionTriviallyDead(cast<Instruction>(U));
+        }) > 1) {
+      LLVM_DEBUG(dbgs() << "Multiply has more than one use\n");
+      return false;
+    }
+
     // Look through extends if the IV has been widened. Don't look through
     // extends if we already looked through a trunc.
     if (Widened && IsAdd &&

diff  --git a/llvm/test/Transforms/LoopFlatten/pr59339.ll b/llvm/test/Transforms/LoopFlatten/pr59339.ll
index 79150fcd8a62..ef78b385e988 100644
--- a/llvm/test/Transforms/LoopFlatten/pr59339.ll
+++ b/llvm/test/Transforms/LoopFlatten/pr59339.ll
@@ -1,14 +1,11 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -S -opaque-pointers -passes='loop(loop-flatten),verify' -verify-loop-info -verify-dom-info -verify-scev | FileCheck %s
 
-; FIXME: This incorrectly triggers, thus resulting in a miscompilation.
-; The problem is that there is a linear use of the induction variable, i*3+j,
-; but only the part i*3 is used for indexing.
+; Check that the mul does not have extra uses.
 
 define void @test0(ptr %arg, ptr %arg1) {
 ; CHECK-LABEL: @test0(
 ; CHECK-NEXT:  bb:
-; CHECK-NEXT:    [[FLATTEN_TRIPCOUNT:%.*]] = mul i64 3, 2
 ; CHECK-NEXT:    br label [[DOTPREHEADER:%.*]]
 ; CHECK:       .preheader:
 ; CHECK-NEXT:    [[I:%.*]] = phi i64 [ 0, [[BB:%.*]] ], [ [[I5:%.*]], [[BB4:%.*]] ]
@@ -17,17 +14,17 @@ define void @test0(ptr %arg, ptr %arg1) {
 ; CHECK-NEXT:    br label [[BB7:%.*]]
 ; CHECK:       bb4:
 ; CHECK-NEXT:    [[I5]] = add i64 [[I]], 1
-; CHECK-NEXT:    [[I6:%.*]] = icmp slt i64 [[I5]], [[FLATTEN_TRIPCOUNT]]
+; CHECK-NEXT:    [[I6:%.*]] = icmp slt i64 [[I5]], 2
 ; CHECK-NEXT:    br i1 [[I6]], label [[DOTPREHEADER]], label [[BB14:%.*]]
 ; CHECK:       bb7:
-; CHECK-NEXT:    [[I8:%.*]] = phi i64 [ 0, [[DOTPREHEADER]] ]
+; CHECK-NEXT:    [[I8:%.*]] = phi i64 [ 0, [[DOTPREHEADER]] ], [ [[I12:%.*]], [[BB7]] ]
 ; CHECK-NEXT:    [[I9:%.*]] = load i16, ptr [[I3]], align 2
 ; CHECK-NEXT:    [[I10:%.*]] = add i64 [[I8]], [[I2]]
-; CHECK-NEXT:    [[I11:%.*]] = getelementptr i16, ptr [[ARG:%.*]], i64 [[I]]
+; CHECK-NEXT:    [[I11:%.*]] = getelementptr i16, ptr [[ARG:%.*]], i64 [[I10]]
 ; CHECK-NEXT:    store i16 [[I9]], ptr [[I11]], align 2
-; CHECK-NEXT:    [[I12:%.*]] = add nuw nsw i64 [[I8]], 1
+; CHECK-NEXT:    [[I12]] = add nuw nsw i64 [[I8]], 1
 ; CHECK-NEXT:    [[I13:%.*]] = icmp ult i64 [[I12]], 3
-; CHECK-NEXT:    br label [[BB4]]
+; CHECK-NEXT:    br i1 [[I13]], label [[BB7]], label [[BB4]]
 ; CHECK:       bb14:
 ; CHECK-NEXT:    ret void
 ;