Support bufferizable in okl

oneflow/core/job/intra_job_mem_sharing_util.cpp

@@ -56,11 +56,6 @@ namespace oneflow {
 
 namespace {
 
-struct MemBlockResultInfo {
-  size_t mem_block_size;
-  HashMap<RegstDescProto*, int64_t> regst_desc2offset;
-};
-
 int64_t GenDeviceUniqueId(int64_t machine_id, int64_t device_id) {
   return (machine_id << 32) | device_id;
 }
@@ -411,154 +406,11 @@ void GenRegstAllocFreeTimeLineAndRegstLifetimes(
   CHECK(remain_regsts.empty());
 }
 
-// Judge whether a is suitable than b for a gap
-bool SuitableThan(int64_t a, int64_t b) {
-  // The number have orders
-  // A non-negative number is always more suitable than a negative number
-  // If a number is non-negative, then the smaller the better
-  // If a number is negative, then the larger the better
-  // 0 > 1 > 2 > ... > 999999999 > -1 > -2 > ... > -99999999
-  // Now we flip the positive part to make it "the larger the better".
-  if (a >= 0) { a = GetMaxVal<int64_t>() - a; }
-  if (b >= 0) { b = GetMaxVal<int64_t>() - b; }
-  return a > b;
-}
-
-void MemReusedAlgorithmAllocateByOrder(
-    const bool compact_insert, const std::vector<RegstDescProto*>& order,
-    const HashMap<RegstDescProto*, size_t>& regst_desc2size,
-    const HashMap<RegstDescProto*, std::pair<int32_t, int32_t>>& regst2lifetime,
-    MemBlockResultInfo* result) {
-  HashMap<RegstDescProto*, int64_t>* regst_desc2offset = &(result->regst_desc2offset);
-  // NOTE: It is important to make the variables local.
-  // It took me several days to find out that using passed-in vector for size, order, and lifetime
-  // would double the running time. Switch HashMap to vector
-  int32_t total_register_num = order.size();
-  std::vector<int64_t> order2size(total_register_num);
-  std::vector<std::pair<int32_t, int32_t>> order2lifetime(total_register_num);
-  std::vector<int64_t> order2offset(total_register_num);
-  for (int32_t i = 0; i < total_register_num; i++) {
-    order2size[i] = regst_desc2size.at(order[i]);
-    order2lifetime[i] = regst2lifetime.at(order[i]);
-  }
-  size_t buffer_size = 1;
-  // Sort by offset
-  auto comp = [&order2offset](const auto& a, const auto& b) {
-    if (order2offset[a] != order2offset[b]) { return order2offset[a] < order2offset[b]; }
-    // Make sure we have a stable order even if we have the same offset for different registers
-    return a < b;
-  };
-  std::set<int32_t, decltype(comp)> sorted_registers(comp);
-  // Decide offset following the given order
-  for (int32_t inserting_id = 0; inserting_id < total_register_num; inserting_id++) {
-    const auto& inserting_lifetime = order2lifetime[inserting_id];
-    // At the beginning, try to insert the offset in the front of the whole memory pool.
-    int64_t inserting_offset = 0;
-    int64_t inserting_end = inserting_offset + order2size[inserting_id];
-    if (compact_insert) {
-      // Find the most suitable gap for the register
-      int64_t gap_head = 0;
-      int64_t inserting_size = order2size[inserting_id];
-      // difference = length of gap - length of the inserting register
-      int64_t diff_gap = 0, suitable_diff_gap = -1 - inserting_size;
-      for (const auto& curr_register : sorted_registers) {
-        // Ignore those non-excluded registers
-        if (IsLifetimeExcluded(inserting_lifetime, order2lifetime[curr_register])) {
-          if (gap_head < order2offset[curr_register]) {
-            // Find one gap
-            diff_gap = (order2offset[curr_register] - gap_head) - inserting_size;
-            // Compared with the previous suitable gap
-            if (SuitableThan(diff_gap, suitable_diff_gap)) {
-              suitable_diff_gap = diff_gap;
-              // We may insert the register into the gap
-              inserting_offset = gap_head;
-            }
-            // Update gap head
-            gap_head = order2offset[curr_register] + order2size[curr_register];
-          } else {
-            // No gap, update gap head
-            gap_head = std::max(gap_head, order2offset[curr_register] + order2size[curr_register]);
-          }
-        }
-      }
-      // Deal with the buffer_size, which may be the final gap
-      diff_gap = (buffer_size - gap_head) - inserting_size;
-      // Compared with the previous suitable gap
-      if (SuitableThan(diff_gap, suitable_diff_gap)) {
-        suitable_diff_gap = diff_gap;
-        // We may insert the register into the gap
-        inserting_offset = gap_head;
-      }
-      // If no gap large enough to contain the current register
-      if (suitable_diff_gap < 0) {
-        // Prolong the maximum memory pool size by (-suitable_diff_gap)
-        buffer_size -= suitable_diff_gap;
-        int64_t gap_end = suitable_diff_gap + inserting_size + inserting_offset;
-        for (auto reverse_it = sorted_registers.rbegin(); reverse_it != sorted_registers.rend();
-             reverse_it++) {
-          // All the registers with offset < gap_end maintain their position
-          if (order2offset[*reverse_it] < gap_end) { break; }
-          // All the registers with offset >= gap_end move backward
-          order2offset[*reverse_it] -= suitable_diff_gap;
-        }
-      }
-
-    } else {
-      for (const auto& curr_register : sorted_registers) {
-        // i: inserting register, j: current register
-        // x: register offset, l: register size
-        // If x_i + l_i <= x_j, then the inserting register would be placed at x_i
-        if (order2offset[curr_register] >= inserting_end) { break; }
-        // If i and j are excluded, and x_i + l_i > x_j,
-        // then we try to place i at x_j + l_j and check the following registers
-        if (IsLifetimeExcluded(inserting_lifetime, order2lifetime[curr_register])) {
-          int64_t curr_end = order2offset[curr_register] + order2size[curr_register];
-          // Can not set inserting offset = current end directly.
-          // We might have two excluded registers like this:
-          // register a: [100, 10000]
-          // register b: [500, 600]
-          if (inserting_offset < curr_end) {
-            inserting_offset = curr_end;
-            inserting_end = inserting_offset + order2size[inserting_id];
-          }
-        }
-      }
-      // Update total size
-      if (inserting_end > buffer_size) { buffer_size = inserting_end; }
-    }
-    // Either we break the loop or the loop terminated naturally, we can place i at inserting_offset
-    order2offset[inserting_id] = inserting_offset;
-    sorted_registers.insert(inserting_id);
-  }
-
-  result->mem_block_size = buffer_size;
-  // Switch vector to HashMap
-  for (int32_t i = 0; i < total_register_num; i++) {
-    (*regst_desc2offset)[order[i]] = order2offset[i];
-  }
-}
-
-void MemReusedMemSizeFirstAlgo(
-    const bool compact_insert,
-    const HashMap<RegstDescProto*, std::pair<int32_t, int32_t>>& regst2lifetime,
-    const HashMap<RegstDescProto*, size_t>& mem_reused_regst2size, MemBlockResultInfo* result) {
-  std::vector<RegstDescProto*> order;
-  order.reserve(regst2lifetime.size());
-  for (const auto& pair : regst2lifetime) { order.emplace_back(pair.first); }
-  std::sort(order.begin(), order.end(), [&](RegstDescProto* lhs, RegstDescProto* rhs) {
-    size_t l_value = mem_reused_regst2size.at(lhs);
-    size_t r_value = mem_reused_regst2size.at(rhs);
-    if (l_value == r_value) { return regst2lifetime.at(lhs).first < regst2lifetime.at(rhs).first; }
-    return l_value > r_value;
-  });
-  MemReusedAlgorithmAllocateByOrder(compact_insert, order, mem_reused_regst2size, regst2lifetime,
-                                    result);
-}
-
 void MemReusedLifetimeFirstAlgo(
     const bool compact_insert,
     const HashMap<RegstDescProto*, std::pair<int32_t, int32_t>>& regst2lifetime,
-    const HashMap<RegstDescProto*, size_t>& mem_reused_regst2size, MemBlockResultInfo* result) {
+    const HashMap<RegstDescProto*, size_t>& mem_reused_regst2size,
+    MemBlockResultInfo<RegstDescProto*>* result) {
   std::vector<RegstDescProto*> order;
   order.reserve(regst2lifetime.size());
   for (const auto& pair : regst2lifetime) { order.emplace_back(pair.first); }
@@ -575,7 +427,8 @@ void MemReusedLifetimeFirstAlgo(
 void MemReusedTimeLineAlgo(
     const bool compact_insert,
     const HashMap<RegstDescProto*, std::pair<int32_t, int32_t>>& regst2lifetime,
-    const HashMap<RegstDescProto*, size_t>& mem_reused_regst2size, MemBlockResultInfo* result) {
+    const HashMap<RegstDescProto*, size_t>& mem_reused_regst2size,
+    MemBlockResultInfo<RegstDescProto*>* result) {
   std::vector<RegstDescProto*> order;
   order.reserve(regst2lifetime.size());
   for (const auto& pair : regst2lifetime) { order.emplace_back(pair.first); }
@@ -594,7 +447,8 @@ void MemReusedTimeLineAlgo(
 void MemReusedMemVolumeFirstAlgo(
     const bool compact_insert,
     const HashMap<RegstDescProto*, std::pair<int32_t, int32_t>>& regst2lifetime,
-    const HashMap<RegstDescProto*, size_t>& mem_reused_regst2size, MemBlockResultInfo* result) {
+    const HashMap<RegstDescProto*, size_t>& mem_reused_regst2size,
+    MemBlockResultInfo<RegstDescProto*>* result) {
   std::vector<RegstDescProto*> order;
   order.reserve(regst2lifetime.size());
   auto ComputeMemoryVolume = [&](RegstDescProto* key) {
@@ -617,7 +471,8 @@ void MemReusedMemVolumeFirstAlgo(
 void SelectAlgorithmGenMemBlockOffset4Regsts(
     MemAllocAlgoType algo_id, const bool compact_insert,
     const HashMap<RegstDescProto*, std::pair<int32_t, int32_t>>& regst2lifetime,
-    const HashMap<RegstDescProto*, size_t>& mem_reused_regst2size, MemBlockResultInfo* result) {
+    const HashMap<RegstDescProto*, size_t>& mem_reused_regst2size,
+    MemBlockResultInfo<RegstDescProto*>* result) {
   CHECK_EQ(result->mem_block_size, 0);
   CHECK(result->regst_desc2offset.empty());
 
@@ -661,7 +516,8 @@ int64_t CountMemAllocAlgoNum() {
   return alloc_algo_num * compact_insert_num;
 }
 
-void InitAlgo2Result(HashMap<std::pair<MemAllocAlgoType, bool>, MemBlockResultInfo>* algo2result) {
+void InitAlgo2Result(
+    HashMap<std::pair<MemAllocAlgoType, bool>, MemBlockResultInfo<RegstDescProto*>>* algo2result) {
   CHECK(algo2result->empty());
   std::vector<bool> compact_insert_algorithms;
   const MemoryCompactInsertConf& mem_compact_insert_conf =
@@ -676,16 +532,16 @@ void InitAlgo2Result(HashMap<std::pair<MemAllocAlgoType, bool>, MemBlockResultIn
   // NOTE: Experiments show that memory first might be good enough for some cases.
   for (auto compact_insert : compact_insert_algorithms) {
     if (mem_alloc_algo_conf.use_mem_size_first_algo()) {
-      (*algo2result)[{kMemSizeFirstAlgo, compact_insert}] = MemBlockResultInfo();
+      (*algo2result)[{kMemSizeFirstAlgo, compact_insert}] = MemBlockResultInfo<RegstDescProto*>();
     }
     if (mem_alloc_algo_conf.use_lifetime_first_algo()) {
-      (*algo2result)[{kLifetimeFirstAlgo, compact_insert}] = MemBlockResultInfo();
+      (*algo2result)[{kLifetimeFirstAlgo, compact_insert}] = MemBlockResultInfo<RegstDescProto*>();
     }
     if (mem_alloc_algo_conf.use_time_line_algo()) {
-      (*algo2result)[{kTimeLineAlgo, compact_insert}] = MemBlockResultInfo();
+      (*algo2result)[{kTimeLineAlgo, compact_insert}] = MemBlockResultInfo<RegstDescProto*>();
     }
     if (mem_alloc_algo_conf.use_mem_volume_first_algo()) {
-      (*algo2result)[{kMemVolumeFirstAlgo, compact_insert}] = MemBlockResultInfo();
+      (*algo2result)[{kMemVolumeFirstAlgo, compact_insert}] = MemBlockResultInfo<RegstDescProto*>();
     }
   }
 }
@@ -725,7 +581,7 @@ void IntraJobMemSharingUtil::InferMemBlockId4MemReusedRegst(
   }
 
   // step 2: multi-thread run several algorithm for each mem chain
-  HashMap<int64_t, HashMap<std::pair<MemAllocAlgoType, bool>, MemBlockResultInfo>>
+  HashMap<int64_t, HashMap<std::pair<MemAllocAlgoType, bool>, MemBlockResultInfo<RegstDescProto*>>>
       mem_chain2algo2result;
   {
     int64_t work_size = mem_chain2mem_reused_regsts.size() * CountMemAllocAlgoNum();
@@ -737,7 +593,7 @@ void IntraJobMemSharingUtil::InferMemBlockId4MemReusedRegst(
       for (auto& pair : mem_chain2algo2result.at(mem_chain_id)) {
         MemAllocAlgoType algo_id = pair.first.first;
         bool compact_insert = pair.first.second;
-        MemBlockResultInfo* result = &pair.second;
+        MemBlockResultInfo<RegstDescProto*>* result = &pair.second;
         thread_pool.AddWork([algo_id, compact_insert, mem_chain_id, &mem_chain2regst2lifetime,
                              &mem_reused_regst2size, result, &counter]() {
           SelectAlgorithmGenMemBlockOffset4Regsts(algo_id, compact_insert,
@@ -752,7 +608,7 @@ void IntraJobMemSharingUtil::InferMemBlockId4MemReusedRegst(
 
   // step 3: choose best one for each mem chain and set offset for inplace consumer regst
   for (auto& pair : mem_chain2algo2result) {
-    MemBlockResultInfo* best_result = nullptr;
+    MemBlockResultInfo<RegstDescProto*>* best_result = nullptr;
     for (auto& algo_result_pair : pair.second) {
       if (!best_result || algo_result_pair.second.mem_block_size < best_result->mem_block_size) {
         best_result = &algo_result_pair.second;