[code generation] bugfix in single buffering

lib/codegen/allocation.cpp

@@ -29,7 +29,7 @@ void allocation::run(){
   std::vector<ir::value *> J = I;
 
   triples_map_type H;
-  H.insert({0, segment{0, 100}});
+  H.insert({0, segment{0, 1024}});
 
   std::vector<ir::value *> V;
   std::map<ir::value *, unsigned> starts;
@@ -116,6 +116,9 @@ void allocation::run(){
   for(auto &x: offsets_){
     allocated_size_ = std::max<size_t>(allocated_size_, x.second + get_num_bytes(x.first));
   }
+  std::cout << "Allocated: " << allocated_size_ << std::endl;
+  for(auto &x: offsets_)
+    std::cout << x.first->get_name() << " " << x.second << std::endl;
 }
 
 }

lib/codegen/barriers.cpp

@@ -6,6 +6,7 @@
 #include "triton/ir/function.h"
 #include "triton/ir/basic_block.h"
 #include "triton/ir/instructions.h"
+#include "triton/ir/cfg.h"
 
 namespace triton {
 
@@ -62,27 +63,76 @@ void barriers::insert_barrier(ir::instruction *instr, ir::builder &builder) {
   }
 }
 
-void barriers::add(ir::basic_block *block, interval_vec_t &not_synced, ir::builder &builder) {
+barriers::interval_vec_t barriers::join(const std::vector<interval_vec_t>& intervals) {
+  barriers::interval_vec_t result;
+  for(auto x: intervals)
+    for(interval_t i: x)
+      result.push_back(i);
+  return result;
+}
+
+std::pair<barriers::interval_vec_t,
+          barriers::interval_vec_t> barriers::transfer(ir::basic_block *block,
+                                            const interval_vec_t &written_to,
+                                            const interval_vec_t &read_from,
+                                            std::set<ir::instruction*>& insert_loc) {
   ir::basic_block::inst_list_t instructions = block->get_inst_list();
+  interval_vec_t new_written_to = written_to;
+  interval_vec_t new_read_from = read_from;
   for(ir::instruction *i: instructions){
     interval_vec_t read, written;
     get_read_intervals(i, read);
     get_written_intervals(i, written);
-    if(intersect(not_synced, read)) {
-      not_synced.clear();
-      insert_barrier(i, builder);
+    bool read_while_written = intersect(new_written_to, read);
+    bool written_while_read = intersect(new_read_from, written);
+    // double buffering: write and phi-node read won't intersect
+    if(dynamic_cast<ir::copy_to_shared_inst*>(i) &&
+       buffer_info_->is_double(buffer_info_->get_reference(i)))
+      written_while_read = false;
+    if(read_while_written || written_while_read) {
+      insert_loc.insert(i);
+      new_written_to.clear();
+      new_read_from.clear();
     }
-    std::copy(written.begin(), written.end(), std::back_inserter(not_synced));
+    std::copy(written.begin(), written.end(), std::back_inserter(new_written_to));
+    std::copy(read.begin(), read.end(), std::back_inserter(new_read_from));
   }
+  return std::make_pair(new_written_to, new_read_from);
 }
 
 void barriers::run(ir::module &mod) {
   ir::builder &builder = mod.get_builder();
   for(ir::function *fn: mod.get_function_list()){
-    // find barrier location
-    interval_vec_t not_synced;
-    for(ir::basic_block *block: fn->blocks())
-      add(block, not_synced, builder);
+    std::vector<ir::basic_block*> rpo = ir::cfg::reverse_post_order(fn);
+    std::map<ir::basic_block*, interval_vec_t> written_to;
+    std::map<ir::basic_block*, interval_vec_t> read_from;
+    std::set<ir::instruction*> insert_locs;
+    size_t n_inserted_im1 = 0;
+    bool done = false;
+    do{
+      // find barrier location
+      for(ir::basic_block *block: rpo){
+        // written to
+        std::vector<interval_vec_t> pred_written_to;
+        for(ir::basic_block* pred: block->get_predecessors())
+          pred_written_to.push_back(written_to[pred]);
+        // read from
+        std::vector<interval_vec_t> pred_read_from;
+        for(ir::basic_block* pred: block->get_predecessors())
+          pred_read_from.push_back(read_from[pred]);
+        // apply transfer function
+        auto result = transfer(block, join(pred_written_to), join(pred_read_from), insert_locs);
+        written_to[block] = result.first;
+        read_from[block] = result.second;
+      }
+      size_t n_inserted_i = insert_locs.size();
+      done = (n_inserted_im1 == n_inserted_i);
+      n_inserted_im1 = n_inserted_i;
+    }while(!done);
+    for(ir::instruction* i: insert_locs){
+      std::cout << i->get_name() << std::endl;
+      insert_barrier(i, builder);
+    }
   }
 }
 

lib/codegen/buffer_info.cpp

@@ -21,6 +21,16 @@ bool buffer_info_pass::is_loop_latch(ir::phi_node *phi, ir::value *terminator){
     throw std::runtime_error("unreachable");
 }
 
+void buffer_info_pass::replace(ir::value* before, ir::value *after) {
+  shared_.erase(before);
+  shared_.insert(after);
+  if(refs_.find(before) != refs_.end()){
+    ir::value* v = refs_.at(before);
+    refs_.erase(before);
+    refs_.insert({after, v});
+  }
+}
+
 void buffer_info_pass::run(ir::module &mod) {
   // Find which buffers are shared
   for(ir::function *fn: mod.get_function_list())

lib/codegen/liveness.cpp

@@ -11,30 +11,43 @@ namespace codegen{
 
 
 // Entry point
-void liveness::run(ir::module &mod) {
-for(ir::function *fn: mod.get_function_list()){
-  // Assigns index to each instruction
-  slot_index index = 0;
-  for(ir::basic_block *block: fn->blocks())
-  for(ir::instruction *instr: block->get_inst_list()){
-    index += 1;
-    indices_.insert({instr, index});
-  }
-  // Liveness analysis
-  // Creates live intervals
-  for(auto i: indices_){
-    ir::value *v = i.first;
-    if(!info_->is_shared(v) || info_->get_reference(v))
-      continue;
-    unsigned start = i.second;
-    unsigned end = start;
-    for(ir::value *u: v->get_users()){
-      start = std::min(start, indices_.at(u));
-      end = std::max(end, indices_.at(u));
-    }
-    intervals_[v] = segment{start, end};
+inline bool is_shared(ir::value* v) {
+  if(auto x = dynamic_cast<ir::copy_to_shared_inst*>(v))
+    return true;
+  if(auto x = dynamic_cast<ir::phi_node*>(v)){
+    bool res = true;
+    for(unsigned inc = 0; inc < x->get_num_incoming(); inc++)
+      res = res && is_shared(x->get_incoming_value(inc));
+    return res;
   }
+  return false;
 }
+
+void liveness::run(ir::module &mod) {
+  for(ir::function *fn: mod.get_function_list()){
+    // Assigns index to each instruction
+    slot_index index = 0;
+    for(ir::basic_block *block: fn->blocks())
+    for(ir::instruction *instr: block->get_inst_list()){
+      index += 1;
+      indices_.insert({instr, index});
+    }
+    // Liveness analysis
+    // Creates live intervals
+    for(auto i: indices_){
+      ir::value *v = i.first;
+      if(!info_->is_shared(v) || info_->get_reference(v))
+        continue;
+      unsigned start = i.second;
+      unsigned end = start;
+      for(ir::value *u: v->get_users()){
+        start = std::min(start, indices_.at(u));
+        end = std::max(end, indices_.at(u));
+      }
+      intervals_[v] = segment{start, end};
+    }
+    std::cout << "Number of intervals: " << intervals_.size() << std::endl;
+  }
 }
 
 }

lib/codegen/selection.cpp

@@ -748,8 +748,6 @@ void selection::lower_tile_instruction(ir::instruction *ins, llvm::IRBuilder<> &
           indices_t b_idx = {idx[1], builder.getInt32(K)};
           Value *a = TA->get_value(a_idx);
           Value *b = TB->get_value(b_idx);
-//          a = ConstantFP::get(builder.getFloatTy(), 1);
-//          b = ConstantFP::get(builder.getFloatTy(), 1);
           res = builder.CreateCall(f_mul_add, {a, b, res});
         }
         result->set_value(idx, res);
@@ -846,6 +844,7 @@ void selection::run(ir::module &src, Module &dst) {
     // create grids
     init_grids(fn, dst_builder, sh_mem_ptr);
 
+
     // iterate through block
     std::map<ir::basic_block*, BasicBlock*> last_block;
     for(ir::basic_block *block: fn->blocks()) {
@@ -854,10 +853,10 @@ void selection::run(ir::module &src, Module &dst) {
       for(ir::instruction *i: block->get_inst_list()){
         BasicBlock *current = dst_builder.GetInsertBlock();
         bool phi_inserted = (dynamic_cast<ir::phi_node*>(i) || dynamic_cast<ir::merge_inst*>(i)) && !current->empty();
-        if(phi_inserted)
-          dst_builder.SetInsertPoint(&*current->getFirstInsertionPt());
+        if(phi_inserted && current->getFirstNonPHI())
+          dst_builder.SetInsertPoint(&*current->getFirstNonPHI());
         lower_instruction(i, dst_builder);
-        if(phi_inserted)
+        if(phi_inserted && current->getFirstNonPHI())
           dst_builder.SetInsertPoint(current);
         last_block[block] = dst_builder.GetInsertBlock();
       }

lib/codegen/shared_copy.cpp

@@ -28,6 +28,12 @@ void place_shared_copy::run(ir::module &mod) {
   for(ir::instruction *i: block->get_inst_list())
   if(info_->is_shared(i) && !info_->is_double(i))
     add_copy(i, builder);
+
+  for(ir::function *fn: mod.get_function_list())
+  for(ir::basic_block *block: fn->blocks())
+  for(ir::instruction *i: block->get_inst_list())
+    if(auto* cts = dynamic_cast<ir::copy_to_shared_inst*>(i))
+      info_->replace(cts->get_operand(0), cts);
 }
 
 }