staging

src/CMakeLists.txt

@@ -4,8 +4,8 @@ add_library(migraph
     generate.cpp
     program.cpp
     shape.cpp
-    opt/optimize.cpp
     opt/memory_coloring.cpp
+    opt/memory_coloring_impl.cpp
 )
 rocm_clang_tidy_check(migraph)
 target_include_directories(migraph PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>)

src/include/migraph/optimize.hpp → src/include/migraph/memory_coloring.hpp

-#ifndef MIGRAPH_GUARD_RTGLIB_OPTIMIZE_HPP
-#define MIGRAPH_GUARD_RTGLIB_OPTIMIZE_HPP
+#ifndef MIGRAPH_GUARD_RTGLIB_MEMORY_COLORING_HPP
+#define MIGRAPH_GUARD_RTGLIB_MEMORY_COLORING_HPP
 
 #include <string>
 #include <migraph/instruction_ref.hpp>
@@ -7,9 +7,9 @@
 namespace migraph {
 struct program;
 
-struct optimize
+struct memory_coloring
 {
-    std::string name() const { return "optimize"; }
+    std::string name() const { return "memory coloring"; }
     void apply(program& p) const;
 };
     

src/opt/common_header.hpp

@@ -4,4 +4,12 @@
 #include <migraph/instruction.hpp>
 #include <migraph/operators.hpp>
 #include <migraph/iterator_for.hpp>
-#endif 
+#define DEBUG_OPT
+
+#ifdef DEBUG_OPT
+#define DEBUG(s) s
+#else
+#define DEBUG(s)
+#endif // DEBUG_OPT
+
+#endif  // MIGRAPH_GUARD_RTGLIB_COMMON_HEADER_HPP

src/opt/dfs.cpp

-namespace migraph {
-
-    void foo() {
-
-    }
-} // namespace migraph

src/opt/memory_coloring.cpp

-#include "memory_coloring.hpp"
+#include <migraph/memory_coloring.hpp>
+#include "memory_coloring_impl.hpp"
 
 namespace migraph {
-void memory_coloring::run()
+void memory_coloring::apply(program &p) const
 {
-
+    memory_coloring_impl opt(&p);
+    opt.run();
+    
 }
-
-} // namespsace migraph
+} // namespace migraph

src/opt/memory_coloring.hpp

-#ifndef MIGRAPH_GUARD_RTGLIB_MEMORY_COLORING_HPP
-#define MIGRAPH_GUARD_RTGLIB_MEMORY_COLORING_HPP
-#include "common_header.hpp"
-
-namespace migraph {
-struct memory_coloring {
-    explicit memory_coloring(program *p) : p_program(p) {}
-    void run();
-    private:
-    program* p_program;
-};
-
-} // namespace migraph
-#endif

src/opt/memory_coloring_impl.cpp

+#include "memory_coloring_impl.hpp"
+#include <set>
+
+namespace migraph {
+void memory_coloring_impl::run()
+{
+    int num_of_instrs = p_program->get_size();
+    if (num_of_instrs == 0)
+        return;
+
+    DEBUG(dump("before memory coloring"));
+    int cur_points = num_of_instrs * 2;
+    instruction_ref iter = std::prev(p_program->end());
+    instruction_ref begin = p_program->begin();
+    std::vector<T_live_interval> live_intervals;
+    std::vector<instruction_ref> dead_instrs;
+    std::vector<int> orderings;
+    int num_of_lives = 0;
+    std::unordered_map<const instruction*, int> instr2Live;
+    std::set<int> live_set;
+    live_set.clear();
+    instr2Live.clear();
+    live_intervals.resize(num_of_instrs);
+    do {
+        const instruction* p_iter = &(*iter);
+        int def_id = -1;
+        if (instr2Live.find(p_iter) != instr2Live.end()) {
+            def_id = instr2Live[p_iter];
+            bool isLit = isLiteral(iter);
+            if (isAllocate(iter) || isLit) {
+                live_intervals[def_id].begin = cur_points;
+                live_intervals[def_id].result = iter->result;
+                live_intervals[def_id].isLiteral = isLit;
+                orderings.push_back(def_id);
+                live_set.erase(def_id);
+            } 
+        } else if (!isParam(iter) && !isOutline(iter) && !isCheckContext(iter)) {
+            // dead instruction.
+            dead_instrs.push_back(iter);
+        }
+        
+        if (!iter->arguments.empty()) {
+            for (auto&& arg : iter->arguments) {
+                if (isParam(arg) || isOutline(arg)) {
+                    continue;
+                } 
+                const instruction* p_arg = &(*arg);
+                if (isAllocate(arg)) {
+                    assert(def_id != -1);
+                    live_intervals[def_id].addUse(cur_points);
+                    instr2Live[p_arg] = def_id;
+                } else if (instr2Live.find(p_arg) == instr2Live.end()) {
+                    int id = num_of_lives++;
+                    live_intervals[id].id = id;
+                    live_intervals[id].end = cur_points;
+                    live_intervals[id].addUse(cur_points);
+                    instr2Live[p_arg] = id;
+                    live_set.insert(id);
+                } 
+            }
+        }
+        cur_points -= 2;
+        iter = std::prev(iter);
+    } while (iter != begin);
+}
+
+#ifdef DEBUG_OPT
+void memory_coloring_impl::dump(std::string str)
+{
+    std::cout << str << std::endl;
+    std::cout << *p_program << std::endl;
+}
+#endif    
+    
+} // namespace migraph

src/opt/memory_coloring_impl.hpp

+#ifndef MIGRAPH_GUARD_RTGLIB_MEMORY_COLORING_IMPL_HPP
+#define MIGRAPH_GUARD_RTGLIB_MEMORY_COLORING_IMPL_HPP
+#include "common_header.hpp"
+
+namespace migraph {
+struct memory_coloring_impl {
+    explicit memory_coloring_impl(program *p) : p_program(p) {}
+    void run();
+    private:
+    bool isParam(const instruction_ref ins)    { return ins->op.name() == "@param"; }
+    bool isAllocate(const instruction_ref ins) { return ins->op.name() == "hip::allocate"; }
+    bool isOutline(const instruction_ref ins)  { return ins->op.name() == "@outline"; }
+    bool isLiteral(const instruction_ref ins)  { return ins->op.name() == "@literal"; }
+    bool isCheckContext(const instruction_ref ins) { return ins->op.name() == "check_context"; }
+
+#ifdef DEBUG_OPT    
+    void dump(std::string);
+#endif    
+
+    program* p_program;
+};
+
+typedef struct live_interval {
+    explicit live_interval() { init(); }
+    bool isValid()  const    { return (begin != -1) && (end != -1); }
+    void addUse(int use)     { use_points.push_back(use); }
+    void init() { begin = -1; end = -1; id = -1; isLiteral = false; }
+    int begin;
+    int end;
+    int id;
+    std::vector<int> use_points;
+    shape result;
+    bool isLiteral;
+} T_live_interval;
+
+} // namespace migraph
+#endif

src/opt/optimize.cpp

-#include <migraph/optimize.hpp>
-#include "memory_coloring.hpp"
-
-namespace migraph {
-void optimize::apply(program &p) const
-{
-    std::cout << p << std::endl;
-    memory_coloring opt(&p);
-    opt.run();
-    int ins_enum = p.get_size();
-    if (ins_enum == 0)
-        return;
-    instruction_ref iter_ins = std::prev(p.end());
-    instruction_ref first_ins = p.begin();
-    do {
-        iter_ins = std::prev(iter_ins);
-        
-    } while (iter_ins != first_ins);
-    
-}
-} // namespace migraph

src/targets/gpu/target.cpp

 #include <migraph/gpu/target.hpp>
 #include <migraph/gpu/lowering.hpp>
-#include <migraph/optimize.hpp>
+#include <migraph/memory_coloring.hpp>
 #include <migraph/gpu/write_literals.hpp>
 #include <migraph/gpu/context.hpp>
 
@@ -9,7 +9,7 @@ namespace gpu {
 
 std::vector<pass> target::get_passes(migraph::context&) const
 {
-    return {lowering{}, optimize{}, write_literals{}};
+    return {lowering{}, memory_coloring{}, write_literals{}};
 }
 
 std::string target::name() const { return "miopen"; }