add memory coloring

src/include/migraph/operators.hpp

@@ -550,6 +550,31 @@ struct div : binary
     std::string name() const { return "div"; }
 };
 
+struct get_mem_ptr 
+{
+    std::string name() const { return "get_mem_ptr:" + std::to_string(offset); }
+    shape compute_shape(std::vector<shape> inputs) const
+    {
+        return inputs.at(1);
+    }
+    argument compute(context&, shape output_shape, std::vector<argument> args) const {
+        return {output_shape, args.at(0).data() + offset};
+    }
+    std::size_t offset = 0;
+};
+
+struct write_literal
+{
+    std::string name() const { return "write_literal"; }
+    shape compute_shape(std::vector<shape> inputs) const
+    {
+        return inputs.at(2);
+    }
+    argument compute(context&, shape output_shape, std::vector<argument> args) const {
+        assert(false);
+    }
+};        
+
 struct outline
 {
     shape s;

src/opt/memory_coloring_impl.cpp

@@ -6,50 +6,68 @@ void memory_coloring_impl::run()
     build();
     if (num_of_lives != 0) {
         DEBUG(dump("---Before memory coloring---"));
-        DEBUG(dump());
+        DEBUG(dump(p_program));
         // Coloring
         while (!alloc_queue.empty()) {
             T_live_interval* interval = alloc_queue.top();
             allocate(interval);
             alloc_queue.pop();
         }
-        for (int i = 0; i < num_of_lives; ++i)
+        rewrite();
+        DEBUG(verify());
+        for (int i = 0; i < num_of_lives; ++i) {
             free(live_intervals[i]);
         }
+    }
 }
 
 bool memory_coloring_impl::allocate(T_live_interval* interval)
 {
     shape s = interval->result;
-    int size = s.bytes();
+    std::size_t size = s.bytes();
     std::size_t element_size = size / s.elements();
     T_live_range& segment = interval->segment;
     int vn = segment.vn;
     std::priority_queue<T_live_range*, std::vector<T_live_range*>, ordering> conflict_queue;
+    std::unordered_map<long long, T_live_range*> offset2Live;
+    offset2Live.clear();    
 
     if (conflict_table.find(vn) != conflict_table.end()) {
         std::set<int>& vn_set = conflict_table[vn];
         for (auto iter = vn_set.begin(), end = vn_set.end(); iter != end; ++iter) {
             T_live_range* range = live_ranges[*iter];
-            if (range->offset != -1)
+            long long offset = range->offset;
+            if (offset != InvalidOffset) {
                 conflict_queue.push(range);
+                if (offset2Live.find(offset) == offset2Live.end()) {
+                    offset2Live[offset] = range;
+                } else {
+                    T_live_range* prev = offset2Live[offset];
+                    assert(prev->offset == offset);
+                    if (prev->size < range->size)
+                        offset2Live[offset] = range;
+                }
+            }
         }
     }
 
-    int offset = 0;
+    long long offset = 0;
     while (!conflict_queue.empty()) {
         T_live_range* range = conflict_queue.top();
-        int cur_offset = range->offset;
+        long long cur_offset = range->offset;
+        if (offset2Live[cur_offset] == range) {
             if ((cur_offset > offset) && (cur_offset - offset) >= size) {
                 break;
             }
             offset = cur_offset + range->size;
             if ((offset % element_size) != 0)
                 offset += (element_size - (offset % element_size));
+        }
         conflict_queue.pop();
     }
     segment.offset = offset;
     DEBUG(segment.dump());
+    required_bytes = std::max(required_bytes, offset + segment.size);
     return true;
 }
 
@@ -62,9 +80,7 @@ void memory_coloring_impl::build()
     instruction_ref iter = std::prev(p_program->end());
     instruction_ref begin = p_program->begin();
     std::vector<instruction_ref> dead_instrs;
-    std::unordered_map<const instruction*, T_live_interval*> instr2Live;
     std::set<int> live_set;
-    T_live_interval* next_def = nullptr;
     // Build live intervals.
     do {
         const instruction* p_iter = &(*iter);
@@ -80,23 +96,25 @@ void memory_coloring_impl::build()
                 alloc_queue.push(def_interval);
                 range.begin = cur_points;
                 range.size = (iter->result).bytes();
-                next_def = def_interval;
                 live_set.erase(range.vn);
             } 
         } else if (!isParam(iter) && !isOutline(iter) && !isCheckContext(iter)) {
             isDead = true;
         }
+        int tieNdx = getInputTieNdx(iter);
         if (!iter->arguments.empty()) {
+            int cnt = -1;
             for (auto&& arg : iter->arguments) {
+                cnt++;
                 if (isParam(arg) || isOutline(arg)) {
                     if (isOutputParam(arg))
                         isDead = false;
                     continue;
                 }
                 const instruction* p_arg = &(*arg);
-                if (isAllocate(arg)) {
-                    // input is from hip::allocate, def is considered as use
-                    // and coalesce the live intervals.
+                if (cnt == tieNdx) {
+                    // input memory is used as this instruction's output.
+                    // def is considered as use. Coalesce the live intervals.
                     def_interval->addUse(cur_points);
                     instr2Live[p_arg] = def_interval;
                 } else if (instr2Live.find(p_arg) == instr2Live.end()) {
@@ -112,10 +130,6 @@ void memory_coloring_impl::build()
                     live_set.insert(max_value_number);
                     live_intervals[id] = interval;
                     live_ranges[max_value_number] = &(interval->segment);
-                    // Keep track of live intervals that are inactive when
-                    // next_def is enqueued.
-                    if (next_def != nullptr)
-                        next_def->inactive_afters.push_back(interval);
                 } else {
                     T_live_interval* interval = instr2Live[p_arg];
                     interval->addUse(cur_points);
@@ -130,10 +144,57 @@ void memory_coloring_impl::build()
     } while (iter != begin);
 }
 
+void memory_coloring_impl::rewrite()
+{
+    instruction_ref end = p_program->end();
+    instruction_ref scratch_param = end;
+    for (auto ins : iterator_for(*p_program)) {
+        const instruction* p_iter = &(*ins);
+        if (isScratchParam(ins)) {
+            scratch_param = ins;
+            int allocated_bytes = ins->result.bytes();
+            if (allocated_bytes < required_bytes) {
+                std::cout << "required bytes: " << required_bytes << "allocated bytes: " << allocated_bytes << std::endl;
+                throw std::runtime_error("insufficent memory for MIGraph");
+            }
+#ifdef DEBUG_OPT
+            float frac = 1.0 * required_bytes/allocated_bytes*100;
+            std::cout << "memory usage percentage: " << to_string(frac) << "%" << std::endl;
+#endif            
+        }
+        if (instr2Live.find(p_iter) != instr2Live.end()) {
+            T_live_interval* interval = instr2Live[p_iter];
+            if (interval->get_offset() == InvalidOffset) {
+                DEBUG(assert(interval->get_begin() == InvalidOffset));
+                continue;
+            }
+            std::size_t offset = interval->get_offset();
+            if (isAllocate(ins)) {
+                if (scratch_param == end)
+                    throw std::runtime_error("missing scratch parameter");
+                p_program->replace_instruction(ins, get_mem_ptr{offset}, scratch_param, ins->arguments.at(0));
+            } else if (isLiteral(ins)) {
+                if (scratch_param == end)
+                    throw std::runtime_error("missing scratch parameter");
+                auto pre = p_program->add_literal(ins->lit);
+                auto index = p_program->add_literal(offset);
+                p_program->replace_instruction(ins, write_literal{}, scratch_param, index, pre);
+            }
+        }
+    }
+    DEBUG(dump("---After rewrite---"));
+    DEBUG(dump(p_program));
+}
+ 
 #ifdef DEBUG_OPT
 void memory_coloring_impl::dump(std::string str)
 {
     std::cout << str << std::endl;
+    
+}
+
+void memory_coloring_impl::dump(program* p_program)
+{
     std::cout << *p_program << std::endl;
 }
 
@@ -145,7 +206,7 @@ void memory_coloring_impl::dump()
             T_live_interval* interval = live_intervals[i];
             interval->dump();
         }
-        std::cout << "conflict table:" << std::endl;
+        std::cout << "---conflict table---" << std::endl;
         for (int i = 0; i <= max_value_number; ++i) {
             std::cout << " segment:" << i;
             std::cout << " =>";
@@ -158,15 +219,38 @@ void memory_coloring_impl::dump()
     }
 }
 
+void memory_coloring_impl::verify()
+{
+    if (num_of_lives > 0) {
+        for (int i = 0; i < num_of_lives; ++i) {
+            T_live_interval* interval = live_intervals[i];
+            T_live_range& segment = interval->segment;
+            if (segment.offset == InvalidOffset)
+                continue;
+            int vn = segment.vn;
+            if (conflict_table.find(vn) != conflict_table.end()) {
+                std::set<int>& vn_set = conflict_table[vn];
+                for (auto iter = vn_set.begin(), end = vn_set.end(); iter != end; ++iter) {
+                    T_live_range* range = live_ranges[*iter];
+                    if (range->offset == InvalidOffset)
+                        continue;
+                    if (!isDisjoin(*range, segment))
+                        assert(false);
+                }
+            }
+        }
+    }
+}
+
 #define GET_INS_ENUM(x) (((x) >> 1) - 1)
     
 void live_range::dump()
 {
     std::cout << " segment:" << vn;
     std::cout << " [" << GET_INS_ENUM(begin)  << ", " << GET_INS_ENUM(end) << "]";
-    if (offset != -1) {
+    if (offset != InvalidOffset) {
         std::cout << " mem:";
-        std::cout << " [" << offset << "," << offset + size << "]";
+        std::cout << " [" << offset << "," << offset + size - 1 << "]";
     }
     std::cout << std::endl;
 }
@@ -181,13 +265,6 @@ void live_interval::dump()
         std::cout << " " << GET_INS_ENUM(use) << ",";
     }
 
-    if (!inactive_afters.empty()) {
-        std::cout << " inactivate:";
-        for (auto iter = inactive_afters.begin(), end = inactive_afters.end(); iter != end; ++iter) {
-            T_live_interval*& interval = *iter;
-            std::cout << " " << interval->id << ",";
-        }
-    }
     if (isLiteral)
         std::cout << " literal";
     std::cout << " " << result;

src/opt/memory_coloring_impl.hpp

@@ -4,12 +4,14 @@
 
 namespace migraph {
 
+#define InvalidOffset -1
+    
 typedef struct live_range {
     int begin;  // begin point in the instruction stream.
     int end;    // end point in the instruction stream.
-    int offset; // offset to base pointer of allocated memory trunk.
+    long long offset; // offset to base pointer of allocated memory trunk.
     int vn;     // value number that identifies this live_range.
-    int size;   // size of required memory in bytes
+    long long size;   // size of required memory in bytes
 #ifdef DEBUG_OPT
     void dump();
 #endif    
@@ -20,11 +22,12 @@ typedef struct live_interval {
     
     void init() {
         id = -1; isLiteral = false;
-        segment = { -1, -1, -1, -1, 0};
+        segment = { -1, -1, InvalidOffset, -1, 0};
     }
     void addUse(int use)     { use_points.push_front(use); }
     int get_begin() const  { return segment.begin; }
     int get_end()   const  { return segment.end; }
+    long long get_offset() const { return segment.offset; }
 
 #ifdef DEBUG_OPT
     void dump();
@@ -33,9 +36,6 @@ typedef struct live_interval {
     T_live_range segment;
     int id;
     std::list<int> use_points;
-    // Live intervals that are inactive when this live interval is enqueued.
-    // can be used for live interval collapsing.
-    std::list<struct live_interval*> inactive_afters;
     shape result;
     bool isLiteral;
 
@@ -44,11 +44,13 @@ typedef struct live_interval {
 struct memory_coloring_impl {
     explicit memory_coloring_impl(program *p) : p_program(p)
     {
+        instr2Live.clear();
         live_intervals.clear();
         live_ranges.clear();
         conflict_table.clear();
         num_of_lives = 0;
         max_value_number = -1;
+        required_bytes = 0;
     }
     bool allocate(T_live_interval*);
     void addConflicts(std::set<int>& live_set, int val)
@@ -60,20 +62,43 @@ struct memory_coloring_impl {
     }
     void build();
     void run();
+    void rewrite();
     private:
     bool isParam(const instruction_ref ins)    { return ins->op.name() == "@param"; }
     bool isOutputParam(const instruction_ref ins)
     {
         return isParam(ins) && any_cast<builtin::param>(ins->op).parameter == "output";
     }
+    bool isScratchParam(const instruction_ref ins)
+    {
+        return isParam(ins) && any_cast<builtin::param>(ins->op).parameter == "scratch"; 
+    }
     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"; }
+    bool isGPUTranspose(const instruction_ref ins) { return ins->op.name() == "gpu::transpose"; }
+    int getInputTieNdx(const instruction_ref ins) {
+        if (isGPUTranspose(ins))
+            return 0;
+        int cnt = -1;
+        for (auto&& arg : ins->arguments) {
+            cnt++;
+            if (isAllocate(arg))
+                return cnt;
+        }
+        return -1;
+    }
+    bool isDisjoin(T_live_range& range1, T_live_range& range2) {
+        long long end1 = range1.offset + range1.size - 1;
+        long long end2 = range2.offset + range2.size - 1;
+        return  ((end1 < range2.offset) || (end2 < range1.offset));
+    }
     
 #ifdef DEBUG_OPT
     void dump(std::string);
-    void dump();
+    void dump(program*);
+    void verify();
 #endif    
     struct ordering { 
         bool operator() (const T_live_interval* I1, const T_live_interval* I2) const
@@ -94,6 +119,7 @@ struct memory_coloring_impl {
         }
     };
     program* p_program;
+    std::unordered_map<const instruction*, T_live_interval*> instr2Live;
     // Map live interval Id to live interval.
     std::unordered_map<int, T_live_interval*> live_intervals;
     // Map live range value number to live range.
@@ -105,6 +131,7 @@ struct memory_coloring_impl {
     
     int num_of_lives;
     int max_value_number;
+    long long required_bytes;
 };    
 
 } // namespace migraph

src/targets/gpu/hip.cpp

@@ -18,6 +18,8 @@ hip_ptr allocate_gpu(std::size_t sz)
     hipMalloc(&result, sz);
     if (result == nullptr)
         throw std::runtime_error("can not allocate GPU memory");
+    char * ptr = reinterpret_cast<char*>(result);
+    std::cout << "MIGraph allocated mem: [" << result << "," << ptr + sz -1 << "]" << std::endl;
     return hip_ptr{result};
 }
 
@@ -69,6 +71,11 @@ migraph::argument from_gpu(migraph::argument arg)
     return result;
 }
 
+void copy_to_gpu(char* dst, const char* src, std::size_t size)
+{
+    hipMemcpy(dst, src, size, hipMemcpyHostToDevice);
+}
+    
 } // namespace gpu
 
 } // namespace migraph

src/targets/gpu/include/migraph/gpu/hip.hpp

@@ -12,6 +12,8 @@ migraph::argument to_gpu(migraph::argument arg);
 
 migraph::argument from_gpu(migraph::argument arg);
 
+void copy_to_gpu(char* dst, const char* src, std::size_t size);
+
 struct hip_allocate
 {
     std::string name() const { return "hip::allocate"; }
@@ -40,6 +42,23 @@ struct hip_write
     }
 };
 
+struct hip_memcpy
+{
+    std::string name() const { return "hip_memcpy"; }
+    shape compute_shape(std::vector<shape> inputs) const
+    {
+        return inputs.at(2);
+    }
+    argument compute(context&, shape output_shape, std::vector<argument> args) const {
+        std::size_t * p_data = reinterpret_cast<std::size_t*>(args.at(1).data());
+        char* dst = args.at(0).data() + p_data[0];
+        const char* src = args.at(2).data();
+        std::size_t size = args.at(2).get_shape().bytes();
+        copy_to_gpu(dst, src, size);
+        return {output_shape, dst};
+    }
+};        
+
 } // namespace gpu
 
 } // namespace migraph

src/targets/gpu/write_literals.cpp

@@ -11,12 +11,18 @@ void write_literals::apply(program& p) const
 {
     for(auto ins : iterator_for(p))
     {
+#if 0        
         if(ins->op.name() == "@literal")
         {
             literal l = ins->lit;
             auto pre  = p.add_literal(l);
             p.replace_instruction(ins, hip_write{}, pre);
         }
+#else
+        if (ins->op.name() == "write_literal") {
+            p.replace_instruction(ins, hip_memcpy{}, ins->arguments);
+        }
+#endif        
     }
 }