[LoopUnrollAndJam] Visit phi operand dependencies in post-order

[caojoshua]

Summary

Fixes https://github.com/llvm/llvm-project/issues/58565

The previous implementation visits operands in pre-order, but this does
not guarantee an instruction is visited before its uses. This can cause
instructions to be copied in the incorrect order. For example:

a = ...
b = add a, 1
c = add a, b
d = add b, a

Pre-order visits does not guarantee the order in which a and b are
visited. LoopUnrollAndJam may incorrectly insert b before a.

This patch implements post-order visits. By visiting dependencies first,
we guarantee that an instruction's dependencies are visited first.

[caojoshua]

ping

[dmgreen]

Thanks for the patch. It looks OK to me.

llvm/lib/Transforms/Utils/LoopUnrollAndJam.cpp

142

It can be:

std::function<bool (Instruction *I)> ProcessInstr = [&](Instruction *I) {

Then is needn't pass ProcessInstr as a parameter.

177

This needs a bit of formatting.

Alternatively, If I is in AftBlocks I don't think it can be in InsertLocBB too? InsertLocBB should a block before the inner loop.

llvm/test/Transforms/LoopUnrollAndJam/dependencies_visit_order.ll

3

Does it need both run lines?

[caojoshua]

• explicit type of ProcessInstr, do not pass it to itself as parameter
• if I is in AftBlocks, we don't need to check its parent is in LocBB
• remove unnecessary run from test

[caojoshua]

llvm/lib/Transforms/Utils/LoopUnrollAndJam.cpp

142

Yep. Thanks for C++help.

177

Formatted with clang-format. Removed check for InsertLocBB.

llvm/test/Transforms/LoopUnrollAndJam/dependencies_visit_order.ll

3

Nope. Removed second line.

[dmgreen]

Thanks. LGTM

llvm/lib/Transforms/Utils/LoopUnrollAndJam.cpp

 // otherwise returns true. This is used to process the instructions in the
 // Aft blocks that need to be moved before the subloop. It is used in two
 // places. One to check that the required set of instructions can be moved
 // before the loop. Then to collect the instructions to actually move in
 // moveHeaderPhiOperandsToForeBlocks.
 template <typename T>
 static bool processHeaderPhiOperands(BasicBlock *Header, BasicBlock *Latch,
                                      BasicBlockSet &AftBlocks, T Visit) {
-  SmallVector<Instruction *, 8> Worklist;
   SmallPtrSet<Instruction *, 8> VisitedInstr;
-  for (auto &Phi : Header->phis()) {
-    Value *V = Phi.getIncomingValueForBlock(Latch);
-    if (Instruction *I = dyn_cast<Instruction>(V))
-      Worklist.push_back(I);
-  }
 
-  while (!Worklist.empty()) {
-    Instruction *I = Worklist.pop_back_val();
-    if (!Visit(I))
-      return false;
+  auto ProcessInstr = [&](Instruction *I, auto &&ProcessInstr) {
+    if (VisitedInstr.count(I))
+      return true;
+
     VisitedInstr.insert(I);
 
     if (AftBlocks.count(I->getParent()))
       for (auto &U : I->operands())
         if (Instruction *II = dyn_cast<Instruction>(U))
-          if (!VisitedInstr.count(II))
-            Worklist.push_back(II);
+          if (!ProcessInstr(II, ProcessInstr))
+            return false;
+
+    return Visit(I);
+  };
+
+  for (auto &Phi : Header->phis()) {
+    Value *V = Phi.getIncomingValueForBlock(Latch);
+    if (Instruction *I = dyn_cast<Instruction>(V))
+      if (!ProcessInstr(I, ProcessInstr))
+        return false;
   }
 
   return true;
 }
 
 // Move the phi operands of Header from Latch out of AftBlocks to InsertLoc.
 static void moveHeaderPhiOperandsToForeBlocks(BasicBlock *Header,
                                               BasicBlock *Latch,
                                               Instruction *InsertLoc,
                                               BasicBlockSet &AftBlocks) {
   // We need to ensure we move the instructions in the correct order,
   // starting with the earliest required instruction and moving forward.
-  std::vector<Instruction *> Visited;
+  BasicBlock *InsertLocBB = InsertLoc->getParent();
   processHeaderPhiOperands(Header, Latch, AftBlocks,
-                           [&Visited, &AftBlocks](Instruction *I) {
-                             if (AftBlocks.count(I->getParent()))
-                               Visited.push_back(I);
+                           [&AftBlocks, &InsertLoc, &InsertLocBB](Instruction *I) {
+                             if (AftBlocks.count(I->getParent()) && I->getParent() != InsertLocBB)
+                              I->moveBefore(InsertLoc);
                              return true;
                            });
-
-  // Move all instructions in program order to before the InsertLoc
-  BasicBlock *InsertLocBB = InsertLoc->getParent();
-  for (Instruction *I : reverse(Visited)) {
-    if (I->getParent() != InsertLocBB)
-      I->moveBefore(InsertLoc);
-  }
 }
 
 /*
   This method performs Unroll and Jam. For a simple loop like:
   for (i = ..)
     Fore(i)
     for (j = ..)
       SubLoop(i, j)

llvm/test/Transforms/LoopUnrollAndJam/dependencies_visit_order.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes=loop-unroll-and-jam -allow-unroll-and-jam -unroll-and-jam-count=4 < %s -S | FileCheck %s
+; RUN: opt -aa-pipeline=basic-aa -passes='loop-unroll-and-jam' -allow-unroll-and-jam -unroll-and-jam-count=4 < %s -S | FileCheck %s
+
+target datalayout = "e-m:e-p:32:32-i64:64-v128:64:128-a:0:32-n32-S64"
+
+define void @test1() {
+; CHECK-LABEL: @test1(
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    br i1 false, label [[BB1_BB43_CRIT_EDGE_UNR_LCSSA:%.*]], label [[BB_NEW:%.*]]
+; CHECK:       bb.new:
+; CHECK-NEXT:    br label [[BB5_PREHEADER:%.*]]
+; CHECK:       bb5.preheader:
+; CHECK-NEXT:    [[I10:%.*]] = phi i16 [ 0, [[BB_NEW]] ], [ [[I42_3:%.*]], [[BB38:%.*]] ]
+; CHECK-NEXT:    [[NITER:%.*]] = phi i16 [ 0, [[BB_NEW]] ], [ [[NITER_NEXT_3:%.*]], [[BB38]] ]
+; CHECK-NEXT:    [[I4017:%.*]] = zext i16 [[I10]] to i64
+; CHECK-NEXT:    [[I13_I:%.*]] = add nuw nsw i64 [[I4017]], 1
+; CHECK-NEXT:    [[I42:%.*]] = trunc i64 [[I13_I]] to i16
+; CHECK-NEXT:    [[NITER_NEXT:%.*]] = add nuw nsw i16 [[NITER]], 1
+; CHECK-NEXT:    [[I4017_1:%.*]] = zext i16 [[I42]] to i64
+; CHECK-NEXT:    [[I13_I_1:%.*]] = add nuw nsw i64 [[I4017_1]], 1
+; CHECK-NEXT:    [[I42_1:%.*]] = trunc i64 [[I13_I_1]] to i16
+; CHECK-NEXT:    [[NITER_NEXT_1:%.*]] = add nuw nsw i16 [[NITER_NEXT]], 1
+; CHECK-NEXT:    [[I4017_2:%.*]] = zext i16 [[I42_1]] to i64
+; CHECK-NEXT:    [[I13_I_2:%.*]] = add nuw nsw i64 [[I4017_2]], 1
+; CHECK-NEXT:    [[I42_2:%.*]] = trunc i64 [[I13_I_2]] to i16
+; CHECK-NEXT:    [[NITER_NEXT_2:%.*]] = add nuw nsw i16 [[NITER_NEXT_1]], 1
+; CHECK-NEXT:    [[I4017_3:%.*]] = zext i16 [[I42_2]] to i64
+; CHECK-NEXT:    [[I13_I_3:%.*]] = add nuw nsw i64 [[I4017_3]], 1
+; CHECK-NEXT:    [[I42_3]] = trunc i64 [[I13_I_3]] to i16
+; CHECK-NEXT:    [[NITER_NEXT_3]] = add nuw i16 [[NITER_NEXT_2]], 1
+; CHECK-NEXT:    br label [[BB10_PREHEADER:%.*]]
+; CHECK:       bb10.preheader:
+; CHECK-NEXT:    br i1 true, label [[BB38]], label [[BB10_PREHEADER]]
+; CHECK:       bb38:
+; CHECK-NEXT:    [[NITER_NCMP_3:%.*]] = icmp eq i16 [[NITER_NEXT_3]], -28
+; CHECK-NEXT:    br i1 [[NITER_NCMP_3]], label [[BB1_BB43_CRIT_EDGE_UNR_LCSSA_LOOPEXIT:%.*]], label [[BB5_PREHEADER]], !llvm.loop [[LOOP1:![0-9]+]]
+; CHECK:       bb1.bb43_crit_edge.unr-lcssa.loopexit:
+; CHECK-NEXT:    [[I10_UNR_PH:%.*]] = phi i16 [ [[I42_3]], [[BB38]] ]
+; CHECK-NEXT:    br label [[BB1_BB43_CRIT_EDGE_UNR_LCSSA]]
+; CHECK:       bb1.bb43_crit_edge.unr-lcssa:
+; CHECK-NEXT:    [[I10_UNR:%.*]] = phi i16 [ 0, [[BB:%.*]] ], [ [[I10_UNR_PH]], [[BB1_BB43_CRIT_EDGE_UNR_LCSSA_LOOPEXIT]] ]
+; CHECK-NEXT:    br i1 true, label [[BB5_PREHEADER_EPIL_PREHEADER:%.*]], label [[BB1_BB43_CRIT_EDGE:%.*]]
+; CHECK:       bb5.preheader.epil.preheader:
+; CHECK-NEXT:    br label [[BB5_PREHEADER_EPIL:%.*]]
+; CHECK:       bb5.preheader.epil:
+; CHECK-NEXT:    [[I10_EPIL:%.*]] = phi i16 [ [[I10_UNR]], [[BB5_PREHEADER_EPIL_PREHEADER]] ], [ [[I42_EPIL:%.*]], [[BB38_EPIL:%.*]] ]
+; CHECK-NEXT:    [[EPIL_ITER:%.*]] = phi i16 [ 0, [[BB5_PREHEADER_EPIL_PREHEADER]] ], [ [[EPIL_ITER_NEXT:%.*]], [[BB38_EPIL]] ]
+; CHECK-NEXT:    br label [[BB10_PREHEADER_EPIL:%.*]]
+; CHECK:       bb10.preheader.epil:
+; CHECK-NEXT:    br i1 true, label [[BB38_EPIL]], label [[BB10_PREHEADER_EPIL]]
+; CHECK:       bb38.epil:
+; CHECK-NEXT:    [[I4017_EPIL:%.*]] = zext i16 [[I10_EPIL]] to i64
+; CHECK-NEXT:    [[I13_I_EPIL:%.*]] = add i64 [[I4017_EPIL]], 1
+; CHECK-NEXT:    [[I14_I_EPIL:%.*]] = select i1 true, i64 [[I13_I_EPIL]], i64 [[I4017_EPIL]]
+; CHECK-NEXT:    [[I42_EPIL]] = trunc i64 [[I14_I_EPIL]] to i16
+; CHECK-NEXT:    [[I3_NOT_EPIL:%.*]] = icmp eq i16 [[I10_EPIL]], -27
+; CHECK-NEXT:    [[EPIL_ITER_NEXT]] = add i16 [[EPIL_ITER]], 1
+; CHECK-NEXT:    [[EPIL_ITER_CMP:%.*]] = icmp ne i16 [[EPIL_ITER_NEXT]], 2
+; CHECK-NEXT:    br i1 [[EPIL_ITER_CMP]], label [[BB5_PREHEADER_EPIL]], label [[BB1_BB43_CRIT_EDGE_EPILOG_LCSSA:%.*]], !llvm.loop [[LOOP3:![0-9]+]]
+; CHECK:       bb1.bb43_crit_edge.epilog-lcssa:
+; CHECK-NEXT:    br label [[BB1_BB43_CRIT_EDGE]]
+; CHECK:       bb1.bb43_crit_edge:
+; CHECK-NEXT:    ret void
+;
+bb:
+  br label %bb5.preheader
+
+bb5.preheader:                                    ; preds = %bb38, %bb
+  %i10 = phi i16 [ 0, %bb ], [ %i42, %bb38 ]
+  br label %bb10.preheader
+
+bb10.preheader:                                   ; preds = %bb10.preheader, %bb5.preheader
+  br i1 true, label %bb38, label %bb10.preheader
+
+bb38:                                             ; preds = %bb10.preheader
+  %i4017 = zext i16 %i10 to i64
+  %i13.i = add i64 %i4017, 1
+  %i14.i = select i1 true, i64 %i13.i, i64 %i4017
+  %i42 = trunc i64 %i14.i to i16
+  %i3.not = icmp eq i16 %i10, -27
+  br i1 %i3.not, label %bb1.bb43_crit_edge, label %bb5.preheader
+
+bb1.bb43_crit_edge:                               ; preds = %bb38
+  ret void
+}