Merge branch 'mem-color-operand-alias'

src/include/migraph/memory_coloring.hpp

@@ -12,6 +12,7 @@ struct program;
 struct memory_coloring
 {
     std::string allocation_op{};
+    bool verify = false;
     std::string name() const { return "memory coloring"; }
     void apply(program& p) const;
 };

src/opt/memory_coloring.cpp

@@ -8,7 +8,7 @@ void memory_coloring::apply(program& p) const
 {
     if(!enabled(MIGRAPH_DISABLE_MEMORY_COLORING{}))
     {
-        memory_coloring_impl opt(&p, allocation_op);
+        memory_coloring_impl opt(&p, allocation_op, verify);
         opt.run();
     }
 }

src/opt/memory_coloring_impl.cpp

@@ -7,7 +7,6 @@ void memory_coloring_impl::run()
 {
     MIGRAPH_DEBUG(dump("---Before memory coloring---"));
     MIGRAPH_DEBUG(dump_program());
-    register_operand_alias();
     build();
     if(num_of_lives != 0)
     {
@@ -20,7 +19,8 @@ void memory_coloring_impl::run()
             alloc_queue.pop();
         }
         rewrite();
-        MIGRAPH_DEBUG(verify());
+        if(enable_verify)
+            verify();
     }
 }
 
@@ -130,11 +130,8 @@ void memory_coloring_impl::build()
         {
             is_dead = true;
         }
-        int tie_ndx = get_input_tie_ndx(iter);
-        int cnt     = -1;
         for(auto&& arg : iter->inputs())
         {
-            cnt++;
             if(is_param(arg) || is_outline(arg))
             {
                 if(is_output_param(arg))
@@ -145,15 +142,8 @@ void memory_coloring_impl::build()
                 }
                 continue;
             }
-            const instruction* p_arg = &(*arg);
-            if(cnt == tie_ndx && (def_interval != nullptr))
-            {
-                // input memory is used as this instruction's output.
-                // def is considered as use. Coalesce the live intervals.
-                def_interval->add_use(cur_points);
-                instr2_live[p_arg] = def_interval;
-            }
-            else if(instr2_live.find(p_arg) == instr2_live.end())
+            const instruction* p_arg = &(*instruction::get_output_alias(arg));
+            if(instr2_live.find(p_arg) == instr2_live.end())
             {
                 // First time see a use, create a live interval.
                 int id                = num_of_lives++;
@@ -183,23 +173,6 @@ void memory_coloring_impl::build()
     } while(iter != begin);
 }
 
-void memory_coloring_impl::register_operand_alias()
-{
-    operand_alias["hip::allocate"]     = -1;
-    operand_alias["hip::load_literal"] = -1;
-    operand_alias["@outline"]          = -1;
-    operand_alias["check_context"]     = -1;
-    operand_alias["@literal"]          = -1;
-    operand_alias["@param"]            = -1;
-    operand_alias["transpose"]         = 0;
-    operand_alias["flatten"]           = 0;
-    operand_alias["broadcast"]         = 0;
-    operand_alias["identity"]          = 0;
-    operand_alias["reshape"]           = 0;
-    operand_alias["pass"]              = 0;
-    operand_alias["scalar"]            = 0;
-}
-
 void memory_coloring_impl::rewrite()
 {
     std::vector<std::size_t> dims;
@@ -249,37 +222,6 @@ void memory_coloring_impl::rewrite()
     MIGRAPH_DEBUG(dump_program());
 }
 
-#ifdef MIGRAPH_DEBUG_OPT
-
-void memory_coloring_impl::dump(const std::string& str) { std::cout << str << std::endl; }
-
-void memory_coloring_impl::dump_program() { std::cout << *p_program << std::endl; }
-
-void memory_coloring_impl::dump_intervals()
-{
-    if(num_of_lives > 0)
-    {
-        std::cout << "---live intervals ---" << std::endl;
-        for(int i = 0; i < num_of_lives; ++i)
-        {
-            live_interval& interval = live_intervals[i];
-            interval.dump();
-        }
-        std::cout << "---conflict table---" << std::endl;
-        for(int i = 0; i <= max_value_number; ++i)
-        {
-            std::cout << " segment:" << i;
-            std::cout << " =>";
-            std::set<int>& table = conflict_table[i];
-            for(auto& iter : table)
-            {
-                std::cout << (iter) << ",";
-            }
-        }
-        std::cout << std::endl;
-    }
-}
-
 void memory_coloring_impl::verify()
 {
     if(num_of_lives > 0)
@@ -291,7 +233,9 @@ void memory_coloring_impl::verify()
 
             if(segment.begin == invalid_offset)
             {
-                assert(interval.is_live_on_entry);
+                // TODO: This check breaks on the tests
+                // if(!interval.is_live_on_entry)
+                // MIGRAPH_THROW("interval is not live on entry");
                 continue;
             }
 
@@ -309,13 +253,44 @@ void memory_coloring_impl::verify()
                     if(range->offset == invalid_offset)
                         continue;
                     if(!is_disjoin(*range, segment))
-                        assert(false);
+                        MIGRAPH_THROW("range and segment is not disjoined");
                 }
             }
         }
     }
 }
 
+#ifdef MIGRAPH_DEBUG_OPT
+
+void memory_coloring_impl::dump(const std::string& str) { std::cout << str << std::endl; }
+
+void memory_coloring_impl::dump_program() { std::cout << *p_program << std::endl; }
+
+void memory_coloring_impl::dump_intervals()
+{
+    if(num_of_lives > 0)
+    {
+        std::cout << "---live intervals ---" << std::endl;
+        for(int i = 0; i < num_of_lives; ++i)
+        {
+            live_interval& interval = live_intervals[i];
+            interval.dump();
+        }
+        std::cout << "---conflict table---" << std::endl;
+        for(int i = 0; i <= max_value_number; ++i)
+        {
+            std::cout << " segment:" << i;
+            std::cout << " =>";
+            std::set<int>& table = conflict_table[i];
+            for(auto& iter : table)
+            {
+                std::cout << (iter) << ",";
+            }
+        }
+        std::cout << std::endl;
+    }
+}
+
 // map liveness tracking point to instruction enum.
 static int get_ins_enum(int x)
 {

src/opt/memory_coloring_impl.hpp

@@ -52,16 +52,15 @@ using interval_ptr = live_interval*;
 
 struct memory_coloring_impl
 {
-    memory_coloring_impl(program* p, std::string alloc_op)
-        : p_program(p), allocation_op(std::move(alloc_op))
+    memory_coloring_impl(program* p, std::string alloc_op, bool p_verify)
+        : p_program(p), allocation_op(std::move(alloc_op)), enable_verify(p_verify)
     {
         instr2_live.clear();
         live_ranges.clear();
         conflict_table.clear();
-        num_of_lives     = 0;
-        max_value_number = -1;
-        required_bytes   = 0;
-        operand_alias.clear();
+        num_of_lives       = 0;
+        max_value_number   = -1;
+        required_bytes     = 0;
         earliest_end_point = -1;
         latest_end_point   = -1;
         unify_literals     = false;
@@ -77,7 +76,6 @@ struct memory_coloring_impl
     }
     void build();
     void run();
-    void register_operand_alias();
     void rewrite();
 
     private:
@@ -94,31 +92,6 @@ struct memory_coloring_impl
         return ins->name() == "check_context";
     }
 
-    // get operand alias info.  This is a temporary workaround.
-    int get_input_tie_ndx(const instruction_ref ins)
-    {
-        std::string name = ins->name();
-        if(operand_alias.find(name) != operand_alias.end())
-            return operand_alias[name];
-        if(is_allocate(ins))
-        {
-            // This happens to custom allocators.
-            operand_alias[name] = -1;
-            return -1;
-        }
-        int cnt           = -1;
-        int last_allocate = -1;
-        for(auto&& arg : ins->inputs())
-        {
-            cnt++;
-            if(is_allocate(arg) || is_output_param(arg))
-                last_allocate = cnt;
-        }
-        assert(last_allocate != -1);
-        operand_alias[name] = last_allocate;
-        return last_allocate;
-    }
-#ifdef MIGRAPH_DEBUG_OPT
     static bool is_disjoin(live_range& range1, live_range& range2)
     {
         if((range1.size == 0) || (range2.size == 0))
@@ -127,10 +100,11 @@ struct memory_coloring_impl
         long long end2 = range2.offset + range2.size - 1;
         return ((end1 < range2.offset) || (end2 < range1.offset));
     }
+    void verify();
+#ifdef MIGRAPH_DEBUG_OPT
     void dump(const std::string&);
     void dump_program();
     void dump_intervals();
-    void verify();
 #endif
     struct ordering
     {
@@ -166,7 +140,6 @@ struct memory_coloring_impl
     std::unordered_map<int, std::set<int>> conflict_table;
     // Priority queue for coloring.
     std::priority_queue<interval_ptr, std::vector<interval_ptr>, ordering> alloc_queue;
-    std::unordered_map<std::string, int> operand_alias;
 
     int num_of_lives;
     int max_value_number;
@@ -178,6 +151,7 @@ struct memory_coloring_impl
     // Whether to unify literals into coloring.
     bool unify_literals;
     std::string allocation_op{};
+    bool enable_verify;
 };
 
 } // namespace MIGRAPH_INLINE_NS

src/program.cpp

@@ -282,7 +282,7 @@ void program::compile(const target& t, tracer trace)
 {
     assert(this->validate() == impl->instructions.end());
     this->impl->ctx = t.get_context();
-    if(not trace.enabled() or enabled(MIGRAPH_TRACE_COMPILE{}))
+    if(enabled(MIGRAPH_TRACE_COMPILE{}))
         trace = tracer{std::cout};
     trace(*this);
     trace();