[LoopFlattening] Check for extra uses on Mul

[dmgreen]

Summary

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

[SjoerdMeijer]

Thanks a lot for fixing this!

llvm/test/Transforms/LoopFlatten/pr59339.ll

3–4

This comment can now be removed, or perhaps better adjusted.

[nikic]

Please avoid unnamed instructions/blocks in tests, which make future updates a PITA. You can avoid this using opt -passes=instnamer.

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

bool IsAddTrunc =
                                   m_Value(MatchedItCount)));
 
     if (!MatchedItCount)
       return false;
 
     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 &&
         (isa<SExtInst>(MatchedItCount) || isa<ZExtInst>(MatchedItCount))) {
       assert(MatchedItCount->getType() == InnerInductionPHI->getType() &&
              "Unexpected type mismatch in types after widening");
       MatchedItCount = isa<SExtInst>(MatchedItCount)
                            ? dyn_cast<SExtInst>(MatchedItCount)->getOperand(0)

llvm/test/Transforms/LoopFlatten/pr59339.ll

 ; 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:%.*]] ]
 ; CHECK-NEXT:    [[I2:%.*]] = mul i64 [[I]], 3
 ; CHECK-NEXT:    [[I3:%.*]] = getelementptr i16, ptr [[ARG1:%.*]], i64 [[I2]]
 ; 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
 ;
 bb:
   br label %.preheader
 
 .preheader:                                       ; preds = %bb4, %bb
   %i = phi i64 [ 0, %bb ], [ %i5, %bb4 ]