[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 llvm#59339

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

Author: davemgreen

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

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
 ;