From c622619bcb26fce564a646a5112e0754ca739ba0 Mon Sep 17 00:00:00 2001
From: Philippe Tillet <ptillet@g.harvard.edu>
Date: Tue, 10 Sep 2019 00:37:51 -0400
Subject: [PATCH] more progress

---
 lib/codegen/analysis/grid.cc     |  3 ++-
 lib/codegen/transform/reorder.cc | 22 ++++++++++++++++------
 lib/runtime/function.cc          |  2 +-
 3 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/lib/codegen/analysis/grid.cc b/lib/codegen/analysis/grid.cc
index d7b773aaf..89b3d90a8 100644
--- a/lib/codegen/analysis/grid.cc
+++ b/lib/codegen/analysis/grid.cc
@@ -332,10 +332,10 @@ void grids::create_grids(std::vector<ir::value*> &grids,
   std::set<ir::value*> seen;
   std::function<void(ir::value*)> bind_references = [&](ir::value *v)
   {
-    auto order = reorder_->get_order(v);
     // skip
     if(!v->get_type()->is_tile_ty() || !seen.insert(v).second)
       return;
+    auto order = reorder_->get_order(v);
     // recurse
     if(auto *user = dynamic_cast<ir::user*>(v))
       for(ir::value *op: user->ops())
@@ -360,6 +360,7 @@ void grids::create_grids(std::vector<ir::value*> &grids,
   for(auto &ref: references)
     if(std::find(grids.begin(), grids.end(), ref.second) == grids.end())
       grids.push_back(ref.second);
+  std::cout << grids.size() << std::endl;
 }
 
 
diff --git a/lib/codegen/transform/reorder.cc b/lib/codegen/transform/reorder.cc
index c5bc31d59..875faaab1 100644
--- a/lib/codegen/transform/reorder.cc
+++ b/lib/codegen/transform/reorder.cc
@@ -18,7 +18,6 @@ reorder::reorder(analysis::align* align, analysis::meminfo *mem)
   : align_(align), mem_(mem) { }
 
 std::vector<unsigned> reorder::get_order(ir::value* v) {
-  std::cout << v->get_name() << std::endl;
   return order_.at(v);
 }
 
@@ -26,6 +25,20 @@ void reorder::run(ir::module &mod) {
 
   std::set<ir::io_inst*> io;
 
+  std::function<void(ir::value*)> set_order = [&](ir::value *v) -> void {
+    if(order_.find(v) != order_.end())
+      return;
+    if(ir::user* u = dynamic_cast<ir::user*>(v))
+      for(ir::value* op: u->ops())
+        set_order(op);
+    ir::type* ty = v->get_type();
+    if(!ty->is_tile_ty())
+      return;
+    std::vector<unsigned> order(ty->get_tile_shapes().size());
+    std::iota(order.begin(), order.end(), 0);
+    order_[v] = order;
+  };
+
   // initialize work-list
   for(ir::function *fn: mod.get_function_list())
   for(ir::basic_block *block: ir::cfg::reverse_post_order(fn))
@@ -34,10 +47,8 @@ void reorder::run(ir::module &mod) {
       ir::type* ptr_ty = x->get_pointer_operand()->get_type();
       if(ptr_ty->is_tile_ty())
         io.insert(x);
-      std::vector<unsigned> order(ptr_ty->get_tile_shapes().size());
-      std::iota(order.begin(), order.end(), 0);
-      order_[i] = order;
     }
+    set_order(i);
   }
 
   ir::builder &builder = mod.get_builder();
@@ -48,9 +59,8 @@ void reorder::run(ir::module &mod) {
     std::iota(order.begin(), order.end(), 0);
     std::sort(order.begin(), order.end(), [&](unsigned a, unsigned b) { return max_contiguous[a] > max_contiguous[b]; } );
     std::list<ir::instruction*> work_list;
-    if(order != order_[i]){
+    if(order != order_[i])
       work_list.push_back(i);
-    }
     // rematerialize recursively
     while(!work_list.empty()) {
       ir::instruction* current = work_list.back();
diff --git a/lib/runtime/function.cc b/lib/runtime/function.cc
index f2aa8e4db..28f3895a3 100644
--- a/lib/runtime/function.cc
+++ b/lib/runtime/function.cc
@@ -210,9 +210,9 @@ std::unique_ptr<driver::module> function::make_bin(ir::module &module, driver::c
   peephole.run(module);
   dce.run(module);
   alignment_info.run(module);
-  ir::print(module, std::cout);
   if(target->is_gpu())
     shmem_info.run(module);
+  ir::print(module, std::cout);
   reorder.run(module);
   dce.run(module);
   ir::print(module, std::cout);