Add support for individual loop unroll bounds

ir/function.cpp

@@ -9,6 +9,7 @@
 #include "util/unionfind.h"
 #include <algorithm>
 #include <fstream>
+#include <iostream>
 #include <set>
 #include <unordered_set>
 
@@ -587,8 +588,8 @@ static auto getPhiPredecessors(const Function &F) {
   return map;
 }
 
-void Function::unroll(unsigned k) {
-  if (k == 0)
+void Function::unroll(vector<unsigned> unroll_bounds) {
+  if (unroll_bounds.size() == 0)
     return;
 
   LoopAnalysis la(*this);
@@ -608,6 +609,44 @@ void Function::unroll(unsigned k) {
   // computed bottom-up during the post-order traversal below
   unordered_map<BasicBlock*, vector<BasicBlock*>> loop_nodes;
 
+  // grab all value users before duplication so the list is shorter
+  auto users = getUsers();
+  auto phi_preds = getPhiPredecessors(*this);
+
+  vector<BasicBlock *> loop_headers_in_order;
+
+  for (auto *bb : BB_order) { // BB_order maintains the code order
+    auto it = forest.find(bb);
+    if (it != forest.end()) {
+      // bb is a loop header
+      loop_headers_in_order.emplace_back(bb);
+    }
+  }
+
+  // check size of loop unroll bounds is valid
+  if (loop_headers_in_order.size() != unroll_bounds.size()) {
+
+    // check if unroll_bounds has only one element, if so, use that for all loop
+    // headers
+    if (unroll_bounds.size() == 1) {
+      unsigned bound = unroll_bounds.front();
+      unroll_bounds.clear();
+      unroll_bounds.resize(loop_headers_in_order.size(), bound);
+    } else {
+        std::cerr << "Error: Number of loop headers does not match number of unroll bounds.\n"
+          << "Number of loop headers: " << loop_headers_in_order.size() << "\n"
+          << "Number of unroll bounds: " << unroll_bounds.size() << std::endl;
+        exit(1);
+    }
+  }
+
+  // map loop header -> unroll bound
+  unordered_map<BasicBlock *, unsigned> unroll_bound_map;
+  for (size_t i = 0; i < loop_headers_in_order.size(); ++i) {
+    auto *header = loop_headers_in_order[i];
+    unroll_bound_map[header] = unroll_bounds[i];
+  }
+
   // traverse each loop tree in post-order
   while (!worklist.empty()) {
     auto [header, height, flag] = worklist.back();
@@ -672,7 +711,7 @@ void Function::unroll(unsigned k) {
     for (unsigned i = 0; i < height; ++i) {
       name_prefix += "#1";
     }
-    for (unsigned unroll = 2; unroll <= k; ++unroll) {
+    for (unsigned unroll = 2; unroll <= unroll_bound_map[header]; ++unroll) {
       string suffix = name_prefix + '#' + to_string(unroll);
       for (auto *bb : loop_bbs) {
         auto &copies = bbmap.at(bb);
@@ -681,6 +720,10 @@ void Function::unroll(unsigned k) {
       }
     }
 
+    if (unroll_bounds.size() > 1) {
+      unroll_bounds.erase(unroll_bounds.begin());
+    }
+
     // Clone the header once more so that the last iteration of the loop can
     // exit. Otherwise the last iteration would be wasted.
     // Here we assume the header is an exit, as that's the common case.

ir/function.h

@@ -224,7 +224,7 @@ class Function final {
                          const std::vector<std::string_view> &src_glbs);
 
   void topSort();
-  void unroll(unsigned k);
+  void unroll(std::vector<unsigned> unroll_bounds);
 
   void print(std::ostream &os, bool print_header = true) const;
   friend std::ostream &operator<<(std::ostream &os, const Function &f);

llvm_util/cmd_args_def.h

@@ -2,11 +2,11 @@
 // Distributed under the MIT license that can be found in the LICENSE file.
 
 #ifdef ARGS_SRC_TGT
-config::src_unroll_cnt = opt_src_unrolling_factor;
-config::tgt_unroll_cnt = opt_tgt_unrolling_factor;
+config::src_unroll_bounds = std::vector<unsigned>(opt_src_unrolling_factor.begin(), opt_src_unrolling_factor.end());
+config::tgt_unroll_bounds = std::vector<unsigned>(opt_tgt_unrolling_factor.begin(), opt_tgt_unrolling_factor.end());
 #else
-config::src_unroll_cnt = opt_unrolling_factor;
-config::tgt_unroll_cnt = opt_unrolling_factor;
+config::src_unroll_bounds = std::vector<unsigned>({opt_unrolling_factor});
+config::tgt_unroll_bounds = std::vector<unsigned>({opt_unrolling_factor});
 #endif
 config::disable_undef_input = opt_disable_undef;
 config::disable_poison_input = opt_disable_poison;

llvm_util/cmd_args_list.h

@@ -19,13 +19,15 @@ llvm::cl::list<std::string> opt_funcs(LLVM_ARGS_PREFIX "func",
 set<string> func_names;
 
 #ifdef ARGS_SRC_TGT
-llvm::cl::opt<unsigned> opt_src_unrolling_factor(LLVM_ARGS_PREFIX "src-unroll",
-  llvm::cl::desc("Unrolling factor for src function (default=0)"),
-  llvm::cl::init(0), llvm::cl::cat(alive_cmdargs));
+llvm::cl::list<unsigned> opt_src_unrolling_factor(LLVM_ARGS_PREFIX "src-unroll",
+  llvm::cl::desc("Unrolling factors for src function (default=0)"), 
+  llvm::cl::CommaSeparated,
+  llvm::cl::cat(alive_cmdargs));
 
-llvm::cl::opt<unsigned> opt_tgt_unrolling_factor(LLVM_ARGS_PREFIX "tgt-unroll",
-  llvm::cl::desc("Unrolling factor for tgt function (default=0)"),
-  llvm::cl::init(0), llvm::cl::cat(alive_cmdargs));
+llvm::cl::list<unsigned> opt_tgt_unrolling_factor(LLVM_ARGS_PREFIX "tgt-unroll",
+  llvm::cl::desc("Unrolling factors for tgt function (default=0)"),
+  llvm::cl::CommaSeparated,
+  llvm::cl::cat(alive_cmdargs));
 #else
 llvm::cl::opt<unsigned> opt_unrolling_factor(LLVM_ARGS_PREFIX "unroll",
   llvm::cl::desc("Unrolling factor (default=0)"),

tools/transform.cpp

@@ -1976,8 +1976,8 @@ void Transform::preprocess() {
   // bits_program_pointer is used by unroll. Initialize it in advance
   initBitsProgramPointer(*this);
 
-  src.unroll(config::src_unroll_cnt);
-  tgt.unroll(config::tgt_unroll_cnt);
+  src.unroll(config::src_unroll_bounds);
+  tgt.unroll(config::tgt_unroll_bounds);
 }
 
 void Transform::print(ostream &os, const TransformPrintOpts &opt) const {

util/config.cpp

@@ -3,6 +3,7 @@
 
 #include "util/config.h"
 #include <iostream>
+#include <vector>
 
 using namespace std;
 
@@ -20,8 +21,8 @@ bool fail_if_src_is_ub = false;
 bool disallow_ub_exploitation = false;
 bool debug = false;
 bool quiet = false;
-unsigned src_unroll_cnt = 0;
-unsigned tgt_unroll_cnt = 0;
+std::vector<unsigned> src_unroll_bounds;
+std::vector<unsigned> tgt_unroll_bounds;
 unsigned max_offset_bits = 64;
 unsigned max_sizet_bits = 64;
 

util/config.h

@@ -3,8 +3,9 @@
 // Copyright (c) 2018-present The Alive2 Authors.
 // Distributed under the MIT license that can be found in the LICENSE file.
 
-#include <string>
 #include <ostream>
+#include <string>
+#include <vector>
 
 namespace util::config {
 
@@ -33,9 +34,9 @@ extern bool debug;
 
 extern bool quiet;
 
-extern unsigned src_unroll_cnt;
+extern std::vector<unsigned> src_unroll_bounds;
 
-extern unsigned tgt_unroll_cnt;
+extern std::vector<unsigned> tgt_unroll_bounds;
 
 // The maximum number of bits to use for offset computations. Note that this may
 // impact correctness, if values involved in offset computations exceed the