Take loop exit blocks into account in constructing edges for top sort when dealing with a loop header (#908)

Fixes #748.

Author: jkubik101

ir/function.cpp

@@ -59,6 +59,10 @@ void BasicBlock::delInstr(const Instr *i) {
   }
 }
 
+void BasicBlock::addExitBlock(BasicBlock* bb) {
+    exit_blocks.emplace(bb);
+}
+
 vector<Phi*> BasicBlock::phis() const {
   vector<Phi*> phis;
   for (auto &i : m_instrs) {
@@ -358,6 +362,17 @@ static vector<BasicBlock*> top_sort(const vector<BasicBlock*> &bbs) {
       if (dst_I != bb_map.end())
         edges[i].emplace(dst_I->second);
     }
+
+    // If `bb` is a loop header, we need to go through its exit block
+    // in order to account for some transitive dependencies we may have
+    // missed due to compression of its inner loops.
+    // If there are no inner loops, this is redundant and if `bb` is not
+    // a loop header, the set of its exit blocks is empty. 
+    for (auto &dst : bb->getExitBlocks()) {
+      auto dst_I = bb_map.find(dst);
+      if (dst_I != bb_map.end())
+        edges[i].emplace(dst_I->second);
+    }
     ++i;
   }
 
@@ -548,6 +563,7 @@ void Function::unroll(unsigned k) {
       for (auto &dst : bb->targets()) {
         if (!bbmap.count(&dst)) {
           exit_edges.emplace(bb, const_cast<BasicBlock*>(&dst));
+          header->addExitBlock(const_cast<BasicBlock*>(&dst));
         }
       }
     }

ir/function.h

@@ -15,6 +15,7 @@
 #include <string_view>
 #include <tuple>
 #include <unordered_map>
+#include <unordered_set>
 #include <vector>
 
 namespace smt { class Model; }
@@ -27,6 +28,9 @@ class BasicBlock final {
   std::string name;
   std::vector<std::unique_ptr<Instr>> m_instrs;
 
+  // If the basic block is a header, this holds all exit blocks of its loop
+    std::unordered_set<BasicBlock*> exit_blocks; 
+
 public:
   BasicBlock(std::string_view name) : name(name) {}
 
@@ -39,6 +43,11 @@ class BasicBlock final {
   void addInstrAt(std::unique_ptr<Instr> &&i, const Instr *other, bool before);
   void delInstr(const Instr *i);
 
+  void addExitBlock(BasicBlock* bb);
+  const std::unordered_set<BasicBlock*>& getExitBlocks() const {
+    return exit_blocks;
+  }
+
   util::const_strip_unique_ptr<decltype(m_instrs)> instrs() const {
     return m_instrs;
   }

tests/alive-tv/loops/partial_sub_loop_test_branch_loop.srctgt.ll

@@ -0,0 +1,62 @@
+; TEST-ARGS: -src-unroll=2 -tgt-unroll=2
+; llvm/llvm/test/Transforms/LoopSimplifyCFG/constant-fold-branch.ll
+
+define i32 @src(i32 %end) {
+entry:
+  br label %outer_header
+
+outer_header:
+  %j = phi i32 [0, %entry], [%j.inc, %outer_backedge]
+  br label %preheader
+
+preheader:
+  br label %header
+
+header:
+  %i = phi i32 [0, %preheader], [%i.inc, %backedge]
+  br label %backedge
+
+backedge:
+  %i.inc = add i32 %i, 1
+  %cmp = icmp slt i32 %i.inc, %end
+  br i1 %cmp, label %header, label %outer_backedge
+
+outer_backedge:
+  %j.inc = add i32 %j, 1
+  %cmp.j = icmp slt i32 %j.inc, %end
+  br i1 %cmp.j, label %outer_header, label %exit
+
+exit:
+  ret i32 %i.inc
+}
+
+define i32 @tgt(i32 %end) {
+entry:
+  br label %outer_header
+
+outer_header:                                     ; preds = %outer_backedge, %entry
+  %j = phi i32 [ 0, %entry ], [ %j.inc, %outer_backedge ]
+  br label %preheader
+
+preheader:                                        ; preds = %outer_header
+  br label %header
+
+header:                                           ; preds = %backedge, %preheader
+  %i = phi i32 [ 0, %preheader ], [ %i.inc, %backedge ]
+  br label %backedge
+
+backedge:                                         ; preds = %header
+  %i.inc = add i32 %i, 1
+  %cmp = icmp slt i32 %i.inc, %end
+  br i1 %cmp, label %header, label %outer_backedge
+
+outer_backedge:                                   ; preds = %backedge
+  %i.inc.lcssa = phi i32 [ %i.inc, %backedge ]
+  %j.inc = add i32 %j, 1
+  %cmp.j = icmp slt i32 %j.inc, %end
+  br i1 %cmp.j, label %outer_header, label %exit
+
+exit:                                             ; preds = %outer_backedge
+  %i.inc.lcssa.lcssa = phi i32 [ %i.inc.lcssa, %outer_backedge ]
+  ret i32 %i.inc.lcssa.lcssa
+}