Implement loop_info analysis pass

VedantParanjape · VedantParanjape · commit 7b3bc0abc412 · 2023-07-29T12:44:36.000+05:30
This commit adds the analysis pass to find natural loops and subloops.
It uses the algorithm from LLVM LoopInfo anaylsis pass for the same,
heavily uses the dominance tree to find the backedges and loop headers.
diff --git a/include/blocks/dominance.h b/include/blocks/dominance.h
@@ -18,22 +18,26 @@ class dominator_analysis {
         //     std::vector<std::shared_ptr<dominator_tree_>> child_nodes;
         // } dominator_tree;
         dominator_analysis(basic_block::cfg_block &cfg);
+        basic_block::cfg_block &cfg_;
         std::vector<int> &get_postorder_bb_map();
         std::vector<int> &get_postorder();
         std::vector<int> &get_idom();
         std::map<int, std::vector<int>> &get_idom_map();
+        std::vector<int> &get_postorder_idom_map();
         int get_idom(int bb_id);
         std::vector<int> get_idom_map(int bb_id);
+        int get_postorder_idom_map(int idom_id);
         bool dominates(int bb1_id, int bb2_id);
         bool is_reachable_from_entry(int bb_id);
         void analyze();
 
     private:
         std::vector<int> idom;
         std::map<int, std::vector<int>> idom_map;
+        std::vector<int> postorder_idom;
         std::vector<int> postorder;
         std::vector<int> postorder_bb_map;
-        basic_block::cfg_block &cfg_;
+        void postorder_idom_helper(std::vector<bool> &visited, int id);
         void postorder_dfs_helper(std::vector<bool> &visited_bbs, int id);
         void postorder_dfs();
         int intersect(int bb1_id, int bb2_id);
diff --git a/include/blocks/loops.h b/include/blocks/loops.h
@@ -8,9 +8,9 @@
 using namespace block;
 class loop {
 public:
-    loop() = default;
+    loop(std::shared_ptr<basic_block> header): header_block(header) {}
 
-private:
+// private:
     struct loop_bounds_ {
         stmt::Ptr ind_var;
         // MISS: intial value of ind var
@@ -22,17 +22,26 @@ class loop {
         stmt::Ptr entry_stmt;
     } loop_bounds;
 
-    basic_block::cfg_block exit_bbs;
+    basic_block::cfg_block blocks;
+    std::shared_ptr<loop> parent_loop;
+    std::shared_ptr<basic_block> header_block;
+    std::shared_ptr<basic_block> backedge_block;
+    std::vector<std::shared_ptr<loop>> subloops;
 };
 
 class loop_info {
 public:
-    loop_info(basic_block::cfg_block ast, dominator_analysis dt): parent_ast(ast), dta(dt) {
+    loop_info(basic_block::cfg_block ast, dominator_analysis &dt): parent_ast(ast), dta(dt) {
         analyze();
     }
+    std::shared_ptr<loop> allocate_loop(std::shared_ptr<basic_block> header);
+    std::vector<std::shared_ptr<loop>> loops;
+    std::vector<std::shared_ptr<loop>> top_level_loops;
+
 private:
     basic_block::cfg_block parent_ast;
     dominator_analysis dta;
+    std::map<int, std::shared_ptr<loop>> bb_loop_map;
     // discover loops during traversal of the abstract syntax tree
     void analyze();
 };
diff --git a/src/blocks/dominance.cpp b/src/blocks/dominance.cpp
@@ -14,6 +14,17 @@ dominator_analysis::dominator_analysis(basic_block::cfg_block &cfg) : cfg_(cfg)
     analyze();
 }
 
+void dominator_analysis::postorder_idom_helper(std::vector<bool> &visited, int id) {
+    for (int idom_id: idom_map[id]) {
+        std::cerr << idom_id << "\n";
+        if (idom_id != -1 && !visited[idom_id]) {
+            visited[idom_id] = true;
+            postorder_idom_helper(visited, idom_id);
+            postorder_idom.push_back(idom_id);
+        }
+    }
+}
+
 void dominator_analysis::postorder_dfs_helper(std::vector<bool> &visited_bbs, int id) {
     for (auto child: cfg_[id]->successor) {
         if (!visited_bbs[child->id]) {
@@ -48,6 +59,10 @@ std::map<int, std::vector<int>> &dominator_analysis::get_idom_map() {
     return idom_map;
 }
 
+std::vector<int> &dominator_analysis::get_postorder_idom_map() {
+    return postorder_idom;
+}
+
 int dominator_analysis::get_idom(int bb_id) {
     if (bb_id < 0 || bb_id >= (int)idom.size()) {
         return -1;
@@ -64,6 +79,14 @@ std::vector<int> dominator_analysis::get_idom_map(int bb_id) {
     return idom_map[bb_id];
 }
 
+int dominator_analysis::get_postorder_idom_map(int idom_id) {
+    if (idom_id < 0 || idom_id >= (int)postorder_idom.size()) {
+        return -1;
+    }
+
+    return postorder_idom[idom_id];
+}
+
 bool dominator_analysis::dominates(int bb1_id, int bb2_id) {
     if (bb1_id == 0) {
         return true;
@@ -138,11 +161,11 @@ void dominator_analysis::analyze() {
             idom_map[i].push_back(-1);
     }
 
-    // for (auto key: idom_map) {
-    //     std::cout << key.first << ": ";
-    //     for (int id: key.second) {
-    //         std::cout << id << " ";
-    //     }
-    //     std::cout << "\n";
-    // }
+    // build a postorder visit list of idom_tree
+    std::vector<bool> visited_idom_nodes(idom_map.size());
+    visited_idom_nodes.assign(visited_idom_nodes.size(), false);
+    visited_idom_nodes[0] = true;
+
+    postorder_idom_helper(visited_idom_nodes, 0);
+    postorder_idom.push_back(0);
 }
diff --git a/src/blocks/loops.cpp b/src/blocks/loops.cpp
@@ -1,9 +1,114 @@
 #include "blocks/loops.h"
 #include <algorithm>
 
+std::shared_ptr<loop> loop_info::allocate_loop(std::shared_ptr<basic_block> header) {
+    if (!header)
+        return nullptr;
+
+    loops.push_back(std::make_shared<loop>(header));
+    bb_loop_map[header->id] = loops.back();
+    return loops.back();
+}
+
 void loop_info::analyze() {
     std::vector<int> idom = dta.get_idom();
-    for (unsigned int i = 0; i < idom.size(); i++) {
-        std::cout << i << " : " << idom[i] << "\n";
+
+    for (int idom_id: dta.get_postorder_idom_map()) {
+        std::vector<int> backedges;
+        int header = idom_id;
+
+        for (auto backedge: dta.cfg_[header]->predecessor) {
+            if (dta.dominates(header, backedge->id) && dta.is_reachable_from_entry(backedge->id)) {
+                backedges.push_back(backedge->id);
+            }
+        }
+
+        if (!backedges.empty()) {
+            std::shared_ptr<loop> new_loop = allocate_loop(dta.cfg_[header]);
+            if (!new_loop)
+                continue;
+
+            int num_blocks = 0;
+            int num_subloops = 0;
+
+            auto backedge_iter = backedges.begin();
+            // do a reverse CFG traversal to map basic blocks in this loop.
+            basic_block::cfg_block worklist(backedges.size());
+            std::generate(worklist.begin(), worklist.end(), [&backedge_iter, this](){
+                return dta.cfg_[*(backedge_iter++)];
+            });
+
+            while (!worklist.empty()) {
+                unsigned int predecessor_bb_id = worklist.back()->id;
+                worklist.pop_back();
+
+                auto subloop_iter = bb_loop_map.find(predecessor_bb_id);
+                if (subloop_iter == bb_loop_map.end()) {
+                    if (!dta.is_reachable_from_entry(predecessor_bb_id))
+                        continue;
+
+                    bb_loop_map[predecessor_bb_id] = new_loop;
+                    ++num_blocks;
+                    // loop has no blocks between header and backedge
+                    if (predecessor_bb_id == new_loop->header_block->id)
+                        continue;
+
+                    worklist.insert(worklist.end(), dta.cfg_[predecessor_bb_id]->predecessor.begin(), dta.cfg_[predecessor_bb_id]->predecessor.end());
+                }
+                else {
+                    // this block has already been discovered, mapped to some other loop
+                    // find the outermost loop
+                    std::shared_ptr<loop> subloop = subloop_iter->second;
+                    while (subloop->parent_loop) {
+                        subloop = subloop->parent_loop;
+                    }
+
+                    if (subloop == new_loop)
+                        continue;
+
+                    // discovered a subloop of this loop
+                    subloop->parent_loop = new_loop;
+                    ++num_subloops;
+                    num_blocks = num_blocks + subloop->blocks.size();
+                    predecessor_bb_id = subloop->header_block->id;
+
+                    for (auto pred: dta.cfg_[predecessor_bb_id]->predecessor) {
+                        auto loop_iter = bb_loop_map.find(pred->id);
+                        // do not check if loop_iter != bb_loop_map.end(), as a results
+                        // basic blocks that are not directly part of the natural loops
+                        // are skipped, like loop latches.
+                        if (loop_iter->second != subloop)
+                            worklist.push_back(pred);
+                    }
+                }
+            }
+            new_loop->subloops.reserve(num_subloops);
+            new_loop->blocks.reserve(num_blocks);
+        }
     }
-}
+
+    // populate all subloops and loops with blocks
+    for (auto bb_id: dta.get_postorder()) {
+        auto subloop_iter = bb_loop_map.find(bb_id);
+        std::shared_ptr<loop> subloop = nullptr;
+        if (subloop_iter != bb_loop_map.end() && (subloop = subloop_iter->second) && dta.cfg_[bb_id] == subloop_iter->second->header_block) {
+            // check if it is the outermost loop
+            if (subloop->parent_loop != nullptr) {
+                subloop->parent_loop->subloops.push_back(subloop);
+            }
+            else {
+                top_level_loops.push_back(subloop);
+            }
+
+            std::reverse(subloop->blocks.begin(), subloop->blocks.end());
+            std::reverse(subloop->subloops.begin(), subloop->subloops.end());
+
+            subloop = subloop->parent_loop;
+        }
+
+        while (subloop) {
+            subloop->blocks.push_back(dta.cfg_[bb_id]);
+            subloop = subloop->parent_loop;
+        }
+    }
+}
diff --git a/src/builder/builder_context.cpp b/src/builder/builder_context.cpp
@@ -356,9 +356,32 @@ block::stmt::Ptr builder_context::extract_ast_from_function_impl(void) {
 	}
 	std::cerr << "== idom map ==\n";
 
+	std::cerr << "== postorder idom ==\n";
+	for (auto idom: dom.get_postorder_idom_map()) {
+		std::cerr << idom << "\n";
+	}
+	std::cerr << "== postorder idom ==\n";
+
 	std::cerr << "++++++ dominance ++++++ \n";
 
+	std::cerr << "++++++ loop info ++++++ \n";
 	loop_info LI(BBs, dom);
+	int loop_num = 0;
+	for (auto loop: LI.loops) {
+		std::cerr << "++++++ loop " << loop_num++ << " ++++++ \n";
+
+		std::cerr << "loop headers: " << loop->header_block->id << "\n";
+		std::cerr << "blocks: ";
+		for (auto bb: loop->blocks) std::cerr << bb->id << " ";
+		std::cerr << "\n";
+		// std::cerr << "backedge: " << loop->backedge_block->id << "\n";
+		std::cerr << "parent loop: (loop header: " << (loop->parent_loop ? (int)loop->parent_loop->header_block->id : -1) << ")\n";
+		std::cerr << "subloops: ";
+		for (auto subl: loop->subloops) std::cerr << "(loop header: " << subl->header_block->id << ") ";
+		std::cerr << "\n";
+	}
+	std::cerr << "++++++ loop info ++++++ \n";
+
 	if (feature_unstructured)
 		return ast;
 
