[LoopFusion] Move instructions from FC1.Preheader to FC0.Preheader when proven safe.

[Whitney]

Summary

Currently LoopFusion give up when the second loop nest preheader is
not empty. For example:

for (int i = 0; i < 100; ++i) {}
x+=1;
for (int i = 0; i < 100; ++i) {}

The above example should be safe to fuse.
This PR moves instructions in FC1 preheader (e.g. x+=1; ) to
FC0 preheader, which then LoopFusion is able to fuse them.

[Whitney]

Added a new test case.

[jdoerfert]

Some minor comments only.

llvm/lib/Transforms/Scalar/LoopFuse.cpp

89–91

unsafe preheader is probably not the best term. Maybe spell it out: non-empty preheader with instructions that cannot be moved?

llvm/lib/Transforms/Utils/CodeMoverUtils.cpp

485

while (FromBB.size() > 1) might be easier, with *I = FromBB.front().

Why do you take the DT, PDT, DI? I would not or just them to assert(isSafe..)

llvm/test/Transforms/LoopFusion/diagnostics_missed.ll

219

Delete the attributes please.

llvm/test/Transforms/LoopFusion/simple.ll

335

Maybe make it more realistic by adding an argument %x and 1, as your commit message example shows. We can then also have a negative test where you replace 1 with %i from the first loop.

[Whitney]

Addressed review comments.

llvm/lib/Transforms/Utils/CodeMoverUtils.cpp

485

Forgot to add the isSafeToMoveBefore check, there is no need to check for this use case, but this utility can be use by user that didn't check beforehand.

[jdoerfert]

One more minor comment below, otherwise LGTM.

llvm/lib/Transforms/Utils/CodeMoverUtils.cpp

485

Fair.

With the while loop and front, you don't need the terminator check anymore.
For any non-broken IR this will work fine without.

[Whitney]

Address review comment.

[Whitney]

rebase

llvm/include/llvm/Transforms/Utils/CodeMoverUtils.h

bool isControlFlowEquivalent(const BasicBlock &BB0, const BasicBlock &BB1,
                              const DominatorTree &DT,
                              const PostDominatorTree &PDT);
 
 /// Return true if \p I can be safely moved before \p InsertPoint.
 bool isSafeToMoveBefore(Instruction &I, Instruction &InsertPoint,
                         const DominatorTree &DT, const PostDominatorTree &PDT,
                         DependenceInfo &DI);
 
-/// Move instructions from \p FromBB bottom up to the beginning of \p ToBB
-/// when proven safe.
+/// Return true if all instructions (except the terminator) in \p BB can be
+/// safely moved before \p InsertPoint.
+bool isSafeToMoveBefore(BasicBlock &BB, Instruction &InsertPoint,
+                        const DominatorTree &DT, const PostDominatorTree &PDT,
+                        DependenceInfo &DI);
+
+/// Move instructions from \p FromBB bottom up to the beginning of \p ToBB when
+/// proven safe.
 void moveInstsBottomUp(BasicBlock &FromBB, BasicBlock &ToBB,
                        const DominatorTree &DT, const PostDominatorTree &PDT,
                        DependenceInfo &DI);
 
+/// Move instructions from \p FromBB top down to the end of \p ToBB when proven
+/// safe.
+void moveInstsTopDown(BasicBlock &FromBB, BasicBlock &ToBB,
+                      const DominatorTree &DT, const PostDominatorTree &PDT,
+                      DependenceInfo &DI);
+
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_UTILS_CODEMOVERUTILS_H

llvm/lib/Transforms/Scalar/LoopFuse.cpp

 STATISTIC(MayThrowException, "Loop may throw an exception");
 STATISTIC(ContainsVolatileAccess, "Loop contains a volatile access");
 STATISTIC(NotSimplifiedForm, "Loop is not in simplified form");
 STATISTIC(InvalidDependencies, "Dependencies prevent fusion");
 STATISTIC(UnknownTripCount, "Loop has unknown trip count");
 STATISTIC(UncomputableTripCount, "SCEV cannot compute trip count of loop");
 STATISTIC(NonEqualTripCount, "Loop trip counts are not the same");
 STATISTIC(NonAdjacent, "Loops are not adjacent");
-STATISTIC(NonEmptyPreheader, "Loop has a non-empty preheader");
+STATISTIC(
+    UnsafePreheader,
+    "Loop has a non-empty preheader with instructions that cannot be moved");
 STATISTIC(FusionNotBeneficial, "Fusion is not beneficial");
 STATISTIC(NonIdenticalGuards, "Candidates have different guards");
 STATISTIC(NonEmptyExitBlock, "Candidate has a non-empty exit block");
 STATISTIC(NonEmptyGuardBlock, "Candidate has a non-empty guard block");
 STATISTIC(NotRotated, "Candidate is not rotated");
 
 enum FusionDependenceAnalysisChoice {
   FUSION_DEPENDENCE_ANALYSIS_SCEV,
for (auto &CandidateSet : FusionCandidates) {
               !haveIdenticalGuards(*FC0, *FC1)) {
             LLVM_DEBUG(dbgs() << "Fusion candidates do not have identical "
                                  "guards. Not Fusing.\n");
             reportLoopFusion<OptimizationRemarkMissed>(*FC0, *FC1,
                                                        NonIdenticalGuards);
             continue;
           }
 
-          // The following three checks look for empty blocks in FC0 and FC1. If
-          // any of these blocks are non-empty, we do not fuse. This is done
-          // because we currently do not have the safety checks to determine if
-          // it is safe to move the blocks past other blocks in the loop. Once
-          // these checks are added, these conditions can be relaxed.
-          if (!isEmptyPreheader(*FC1)) {
-            LLVM_DEBUG(dbgs() << "Fusion candidate does not have empty "
-                                 "preheader. Not fusing.\n");
+          if (!isSafeToMoveBefore(*FC1->Preheader,
+                                  *FC0->Preheader->getTerminator(), DT, PDT,
+                                  DI)) {
+            LLVM_DEBUG(dbgs() << "Fusion candidate contains unsafe "
+                                 "instructions in preheader. Not fusing.\n");
             reportLoopFusion<OptimizationRemarkMissed>(*FC0, *FC1,
-                                                       NonEmptyPreheader);
+                                                       UnsafePreheader);
             continue;
           }
 
+          // The following two checks look for empty blocks in FC0 and FC1. If
+          // any of these blocks are non-empty, we do not fuse. This is done
+          // because we currently do not have the safety checks to determine if
+          // it is safe to move the blocks past other blocks in the loop. Once
+          // these checks are added, these conditions can be relaxed.
           if (FC0->GuardBranch && !isEmptyExitBlock(*FC0)) {
             LLVM_DEBUG(dbgs() << "Fusion candidate does not have empty exit "
                                  "block. Not fusing.\n");
             reportLoopFusion<OptimizationRemarkMissed>(*FC0, *FC1,
                                                        NonEmptyExitBlock);
             continue;
           }
 
#endif
   /// existing code, which currently has not been implemented.
   Loop *performFusion(const FusionCandidate &FC0, const FusionCandidate &FC1) {
     assert(FC0.isValid() && FC1.isValid() &&
            "Expecting valid fusion candidates");
 
     LLVM_DEBUG(dbgs() << "Fusion Candidate 0: \n"; FC0.dump();
                dbgs() << "Fusion Candidate 1: \n"; FC1.dump(););
 
+    // Move instructions from the preheader of FC1 to the end of the preheader
+    // of FC0.
+    moveInstsTopDown(*FC1.Preheader, *FC0.Preheader, DT, PDT, DI);
+
     // Fusing guarded loops is handled slightly differently than non-guarded
     // loops and has been broken out into a separate method instead of trying to
     // intersperse the logic within a single method.
     if (FC0.GuardBranch)
       return fuseGuardedLoops(FC0, FC1);
 
     assert(FC1.Preheader == FC0.ExitBlock);
     assert(FC1.Preheader->size() == 1 &&

llvm/lib/Transforms/Utils/CodeMoverUtils.cpp

if (std::any_of(InstsToCheck.begin(), InstsToCheck.end(),
                       return true;
                     return false;
                   }))
     return reportInvalidCandidate(I, HasDependences);
 
   return true;
 }
 
+bool llvm::isSafeToMoveBefore(BasicBlock &BB, Instruction &InsertPoint,
+                              const DominatorTree &DT,
+                              const PostDominatorTree &PDT,
+                              DependenceInfo &DI) {
+  return llvm::all_of(BB, [&](Instruction &I) {
+    if (BB.getTerminator() == &I)
+      return true;
+
+    return isSafeToMoveBefore(I, InsertPoint, DT, PDT, DI);
+  });
+}
+
 void llvm::moveInstsBottomUp(BasicBlock &FromBB, BasicBlock &ToBB,
                              const DominatorTree &DT,
                              const PostDominatorTree &PDT, DependenceInfo &DI) {
   for (auto It = ++FromBB.rbegin(); It != FromBB.rend();) {
     Instruction *MovePos = ToBB.getFirstNonPHIOrDbg();
     Instruction &I = *It;
     // Increment the iterator before modifying FromBB.
     ++It;
 
     if (isSafeToMoveBefore(I, *MovePos, DT, PDT, DI))
       I.moveBefore(MovePos);
   }
 }
+
+void llvm::moveInstsTopDown(BasicBlock &FromBB, BasicBlock &ToBB,
+                            const DominatorTree &DT,
+                            const PostDominatorTree &PDT, DependenceInfo &DI) {
+  Instruction *MovePos = ToBB.getTerminator();
+  while (FromBB.size() > 1) {
+    Instruction &I = FromBB.front();
+    if (isSafeToMoveBefore(I, *MovePos, DT, PDT, DI))
+      I.moveBefore(MovePos);
+  }
+}

llvm/test/Transforms/LoopFusion/diagnostics_missed.ll

for.inc13:                                        ; preds = %for.body6
   %inc14 = add nuw nsw i64 %i1.01, 1
   %exitcond = icmp ne i64 %inc14, 200
   br i1 %exitcond, label %for.body6, label %for.cond.cleanup5, !llvm.loop !48
 
 for.end15:                                        ; preds = %for.cond.cleanup5
   ret void
 }
 
-; CHECK: remark: diagnostics_missed.c:38:3: [negative_dependence]: entry and for.end: Loop has a non-empty preheader
+; CHECK: remark: diagnostics_missed.c:38:3: [negative_dependence]: entry and for.end: Dependencies prevent fusion
 define void @negative_dependence(i32* noalias %A) !dbg !51 {
 entry:
   br label %for.body
 
 for.body:                                         ; preds = %entry, %for.inc
   %indvars.iv13 = phi i64 [ 0, %entry ], [ %indvars.iv.next2, %for.inc ]
   %arrayidx = getelementptr inbounds i32, i32* %A, i64 %indvars.iv13
   %tmp = trunc i64 %indvars.iv13 to i32

 
 for.end12:                                        ; preds = %for.inc10
   ret i32 %add, !dbg !76
 }
 
 ; Function Attrs: nounwind readnone speculatable willreturn
 declare void @llvm.dbg.value(metadata, metadata, metadata) #0
 
+; CHECK: remark: diagnostics_missed.c:62:3: [unsafe_preheader]: for.first.preheader and for.second.preheader: Loop has a non-empty preheader with instructions that cannot be moved
+define void @unsafe_preheader(i32* noalias %A, i32* noalias %B) {
+for.first.preheader:
+  br label %for.first, !dbg !80
+
+for.first:
+  %i.02 = phi i64 [ 0, %for.first.preheader ], [ %inc, %for.first ]
+  %arrayidx = getelementptr inbounds i32, i32* %A, i64 %i.02
+  store i32 0, i32* %arrayidx, align 4
+  %inc = add nsw i64 %i.02, 1
+  %cmp = icmp slt i64 %inc, 100
+  br i1 %cmp, label %for.first, label %for.second.preheader
+
+for.second.preheader:
+  call void @bar()
+  br label %for.second
+
+for.second:
+  %j.01 = phi i64 [ 0, %for.second.preheader ], [ %inc6, %for.second ]
+  %arrayidx4 = getelementptr inbounds i32, i32* %B, i64 %j.01
+  store i32 0, i32* %arrayidx4, align 4
+  %inc6 = add nsw i64 %j.01, 1
+  %cmp2 = icmp slt i64 %inc6, 100
+  br i1 %cmp2, label %for.second, label %for.end
+
+for.end:
+  ret void
+}
+declare void @bar()
+
 attributes #0 = { nounwind readnone speculatable willreturn }
 
 !llvm.dbg.cu = !{!2}
 !llvm.module.flags = !{!10, !11, !12, !13}
 
 !0 = !DIGlobalVariableExpression(var: !1, expr: !DIExpression())
 !1 = distinct !DIGlobalVariable(name: "B", scope: !2, file: !3, line: 46, type: !6, isLocal: false, isDefinition: true)
 !2 = distinct !DICompileUnit(language: DW_LANG_C99, file: !3, producer: "clang version 9.0.0 ([e-mail]:compiler/llvm-project.git 23c4baaa9f5b33d2d52eda981d376c6b0a7a3180)", isOptimized: true, runtimeVersion: 0, emissionKind: FullDebug, enums: !4, globals: !5, nameTableKind: GNU)
 !3 = !DIFile(filename: "diagnostics_missed.c", directory: "/tmp")
 !4 = !{}

 !69 = !DILocalVariable(name: "i", scope: !70, file: !3, line: 51, type: !7)
 !70 = distinct !DILexicalBlock(scope: !63, file: !3, line: 51, column: 3)
 !71 = !DILocalVariable(name: "i", scope: !72, file: !3, line: 54, type: !7)
 !72 = distinct !DILexicalBlock(scope: !63, file: !3, line: 54, column: 3)
 !73 = distinct !{!73, !74, !75}
 !74 = !DILocation(line: 51, column: 3, scope: !70)
 !75 = !DILocation(line: 52, column: 15, scope: !70)
 !76 = !DILocation(line: 57, column: 3, scope: !63)
+!77 = distinct !DISubprogram(name: "unsafe_preheader", scope: !3, file: !3, line: 60, type: !15, scopeLine: 60, flags: DIFlagPrototyped, spFlags: DISPFlagDefinition, unit: !2, retainedNodes: !78)
+!78 = !{}
+!79 = distinct !DILexicalBlock(scope: !77, file: !3, line: 3, column: 5)
+!80 = !DILocation(line: 62, column: 3, scope: !79)

llvm/test/Transforms/LoopFusion/guarded.ll

bb9:                                        ; preds = %bb8, %bb9
   br i1 %cmp3, label %bb9, label %bb15
 
 bb15:                               ; preds = %bb9
   br label %bb12
 
 bb12:                                        ; preds = %bb15, %bb14
   ret void
 }
+
+; Test that `%add` is moved in for.first.preheader, and the two loops for.first
+; and for.second are fused.
+
+; CHECK: void @moveinsts_preheader
+; CHECK-LABEL: for.first.guard:
+; CHECK: br i1 %cmp.guard.first, label %for.first.preheader, label %for.end
+; CHECK-LABEL: for.first.preheader:
+; CHECK-NEXT:  %add = add nsw i32 %x, 1
+; CHECK-NEXT:  br label %for.first
+; CHECK-LABEL: for.first:
+; CHECK:   br i1 %cmp.j, label %for.first, label %for.second.exit
+; CHECK-LABEL: for.second.exit:
+; CHECK-NEXT:   br label %for.end
+; CHECK-LABEL: for.end:
+; CHECK-NEXT:   ret void
+define void @moveinsts_preheader(i32* noalias %A, i32* noalias %B, i64 %N, i32 %x) {
+for.first.guard:
+  %cmp.guard.first = icmp slt i64 0, %N
+  br i1 %cmp.guard.first, label %for.first.preheader, label %for.second.guard
+
+for.first.preheader:
+  br label %for.first
+
+for.first:
+  %i = phi i64 [ %inc.i, %for.first ], [ 0, %for.first.preheader ]
+  %Ai = getelementptr inbounds i32, i32* %A, i64 %i
+  store i32 0, i32* %Ai, align 4
+  %inc.i = add nsw i64 %i, 1
+  %cmp.i = icmp slt i64 %inc.i, %N
+  br i1 %cmp.i, label %for.first, label %for.first.exit
+
+for.first.exit:
+  br label %for.second.guard
+
+for.second.guard:
+  %cmp.guard.second = icmp slt i64 0, %N
+  br i1 %cmp.guard.second, label %for.second.preheader, label %for.end
+
+for.second.preheader:
+  %add = add nsw i32 %x, 1
+  br label %for.second
+
+for.second:
+  %j = phi i64 [ %inc.j, %for.second ], [ 0, %for.second.preheader ]
+  %Bj = getelementptr inbounds i32, i32* %B, i64 %j
+  store i32 0, i32* %Bj, align 4
+  %inc.j = add nsw i64 %j, 1
+  %cmp.j = icmp slt i64 %inc.j, %N
+  br i1 %cmp.j, label %for.second, label %for.second.exit
+
+for.second.exit:
+  br label %for.end
+
+for.end:
+  ret void
+}

llvm/test/Transforms/LoopFusion/simple.ll

 for.second.latch:
   %inc.second = add nsw i64 %i.second, 1
   %cmp.second = icmp slt i64 %inc.second, 100
   br i1 %cmp.second, label %for.second, label %for.end
 
 for.end:
   ret void
 }
+
+; Test that `%add` is moved in basic block entry, and the two loops for.first
+; and for.second are fused.
+
+; CHECK: i32 @moveinsts_preheader
+; CHECK-LABEL: entry:
+; CHECK-NEXT: %add = add nsw i32 %x, 1
+; CHECK-NEXT: br label %for.first
+; CHECK-LABEL: for.first:
+; CHECK: br i1 %cmp.j, label %for.first, label %for.second.exit
+; CHECK-LABEL: for.second.exit:
+; CHECK-NEXT: ret i32 %add
+
+define i32 @moveinsts_preheader(i32* %A, i32 %x) {
+entry:
+  br label %for.first
+
+for.first:
+  %i = phi i64 [ 0, %entry ], [ %inc.i, %for.first ]
+  %Ai = getelementptr inbounds i32, i32* %A, i64 %i
+  store i32 0, i32* %Ai, align 4
+  %inc.i = add nsw i64 %i, 1
+  %cmp.i = icmp slt i64 %inc.i, 100
+  br i1 %cmp.i, label %for.first, label %for.first.exit
+
+for.first.exit:
+  %add = add nsw i32 %x, 1
+  br label %for.second
+
+for.second:
+  %j = phi i64 [ 0, %for.first.exit ], [ %inc.j, %for.second ]
+  %Aj = getelementptr inbounds i32, i32* %A, i64 %j
+  store i32 2, i32* %Aj, align 4
+  %inc.j = add nsw i64 %j, 1
+  %cmp.j = icmp slt i64 %inc.j, 100
+  br i1 %cmp.j, label %for.second, label %for.second.exit
+
+for.second.exit:
+  ret i32 %add
+}
+
+; Test that `%add` cannot be moved to basic block entry, as it uses %i, which
+; defined after basic block entry. And the two loops for.first and for.second
+; are not fused.
+
+; CHECK: i64 @unsafe_preheader
+; CHECK-LABEL: entry:
+; CHECK-NEXT: br label %for.first
+; CHECK-LABEL: for.first:
+; CHECK: br i1 %cmp.i, label %for.first, label %for.first.exit
+; CHECK-LABEL: for.first.exit:
+; CHECK-NEXT: %add = add nsw i64 %x, %i
+; CHECK-NEXT: br label %for.second
+; CHECK-LABEL: for.second:
+; CHECK: br i1 %cmp.j, label %for.second, label %for.second.exit
+; CHECK-LABEL: for.second.exit:
+; CHECK-NEXT: ret i64 %add
+
+define i64 @unsafe_preheader(i32* %A, i64 %x) {
+entry:
+  br label %for.first
+
+for.first:
+  %i = phi i64 [ 0, %entry ], [ %inc.i, %for.first ]
+  %Ai = getelementptr inbounds i32, i32* %A, i64 %i
+  store i32 0, i32* %Ai, align 4
+  %inc.i = add nsw i64 %i, 1
+  %cmp.i = icmp slt i64 %inc.i, 100
+  br i1 %cmp.i, label %for.first, label %for.first.exit
+
+for.first.exit:
+  %add = add nsw i64 %x, %i
+  br label %for.second
+
+for.second:
+  %j = phi i64 [ 0, %for.first.exit ], [ %inc.j, %for.second ]
+  %Aj = getelementptr inbounds i32, i32* %A, i64 %j
+  store i32 2, i32* %Aj, align 4
+  %inc.j = add nsw i64 %j, 1
+  %cmp.j = icmp slt i64 %inc.j, 100
+  br i1 %cmp.j, label %for.second, label %for.second.exit
+
+for.second.exit:
+  ret i64 %add
+}