WIP: Implement loop_info pass (works)

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);

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();
 };

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);
 }

src/blocks/loops.cpp

@@ -1,9 +1,133 @@
 #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);
+            }
+        }
+
+        for(int i: backedges) {
+            std::cerr << i << "\n";
+        }
+
+        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++)];
+            });
+            for (unsigned int i = 0; i < worklist.size(); i++) {
+                std::cerr << worklist[i]->id << "\n";
+            }
+            std::cerr << "header-p: " << header << "\n";
+            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;
+                    std::cerr << "bb: " << predecessor_bb_id << "\n";
+                    std::cerr << "loop-h: " << new_loop->header_block->id << "\n";
+                    ++num_blocks;
+                    std::cerr << "if 1\n";
+                    // loop has no blocks between header and backedge
+                    if (predecessor_bb_id == new_loop->header_block->id)
+                        continue;
+                    std::cerr << "if 2\n";
+
+                    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;
+                    std::cerr << "header: " << subloop->header_block->id << "\n";
+                    while (subloop->parent_loop) {
+                        subloop = subloop->parent_loop;
+                    }
+                    std::cerr << "else 1\n";
+
+                    if (subloop == new_loop)
+                        continue;
+
+                    std::cerr << "else 2\n";
+                    std::cerr << "bb: " << predecessor_bb_id << "\n";
+
+                    // 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;
+        std::cerr << bb_id << " " << dta.cfg_[bb_id]->id <<  subloop_iter->second->header_block->id << "\n";
+        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
+            std::cerr << "out::::" << subloop->header_block->id << "\n";
+            if (subloop->parent_loop != nullptr) {
+                std::cerr << "id::::" << subloop->header_block->id << "\n";
+                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;
+        }
+    }
+}

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;