Fix corner case when m is too small

adaptive_schedule.py

@@ -345,10 +345,7 @@ def process_warmup_without_increasing_peak_mem(schedules, m):
 def squeeze_without_change_order(schedules, m):
     p = len(schedules)
     squeezed = [[' '] * len(schedules[_]) for _ in range(p)]
-    max_len = 0
-    for seq in squeezed:
-        assert max_len == 0 or max_len == len(seq)
-        max_len = max(max_len, len(seq))
+    max_len = check_and_get_schedule_len(schedules)
 
     identifier_cnt = [{_id: 0 for _id in "FfBbWw"} for _ in range(p)]
     identifier_index = [{_id: -1 for _id in "FfBbWw"} for _ in range(p * m)]
@@ -389,6 +386,9 @@ def squeeze_without_change_order(schedules, m):
                 identifier_cnt[i][identifier] = _cnt + 1
                 identifier_index[_cnt * p + i][identifier] = index
                 stage_index[i] = index + 1
+    new_len = max(stage_index)
+    for i in range(p):
+        squeezed[i] = squeezed[i][:new_len]
     return squeezed
 
 
@@ -454,6 +454,7 @@ def process_cooldown(schedules, m):
             schedules[i][index] = 'B'
 
     # 2: add W back in cooldown phase
+    max_len = 0
     for i in range(p):
         c_w, c_ww = 0, 0
         last_w_index = -1
@@ -478,12 +479,57 @@ def process_cooldown(schedules, m):
                 elif c_ww > 0:
                     schedules[i][j] = 'w'
                     c_ww -= 1
+        for _ in range(c_w):
+            schedules[i].append('W')
+        for _ in range(c_ww):
+            schedules[i].append('w')
+        max_len = max(max_len, len(schedules[i]))
+    for i in range(p):
+        for _ in range(len(schedules[i]), max_len):
+            schedules[i].append(' ')
 
     schedules = squeeze_without_change_order(schedules, m)
     return schedules
 
 
-def reorder_greedily_without_increasing_peak_mem(schedules, m, starting_index = None, ending_index = None):
+def check_and_get_schedule_len(schedules):
+    max_len = 0
+    for seq in schedules:
+        assert max_len == 0 or max_len == len(seq)
+        max_len = max(max_len, len(seq))
+    return max_len
+
+
+def release_w_in_warmup_if_under_memory(schedules, peak_mem = None):
+    """
+    FF     fBWfBW   bwbw       ->         FF     fBfBWW  bwbw
+     FF   f fBW BW bwbw        ->          FF   f fBWBW bwbw
+      FF f f  BW BbWbww        ->           FF f f  BWBbWbww
+       FfFf    BbWBbwWw        ->            FfFf    BbBbWwWw
+    When the number of micro-batches is too small (than mem), the warmup phase is not optimal. We simply remove some
+    preceding W to fully utilize the memory to reduce unnecessary bubbles.
+    """
+    p = len(schedules)
+    max_len = check_and_get_schedule_len(schedules)
+    all_peak_mem = get_peak_mem(schedules, return_all=True)
+    peak_mem = peak_mem or max(all_peak_mem)
+    min_peak = min(all_peak_mem)
+    for i in range(p):
+        cnt = 0
+        padding = [" "] * (peak_mem - min_peak)
+        for j in range(max_len):
+            if all_peak_mem[i] + cnt >= peak_mem:
+                break
+            if schedules[i][j] in "Ww":
+                padding[cnt] = schedules[i][j]
+                schedules[i][j] = ' '
+                cnt += 1
+        schedules[i].extend(padding)
+    # max_len += peak_mem - min_peak
+    return schedules
+
+
+def reorder_greedily_without_increasing_peak_mem(schedules, m, starting_index = None, ending_index = None, peak_mem = None):
     """
     We iterate all the cells from left to right. If a vacant cell (which means a bubble) is encountered, we try to
     find a computation pass to fill this bubble. We iterate all the following computation passes in the same device,
@@ -491,17 +537,15 @@ def reorder_greedily_without_increasing_peak_mem(schedules, m, starting_index =
     to the vacant cell, and the bubble is filled.
     """
     p = len(schedules)
-    max_len = 0
-    for seq in schedules:
-        assert max_len == 0 or max_len == len(seq)
-        max_len = max(max_len, len(seq))
     if starting_index is not None:
         assert isinstance(starting_index, list) and len(starting_index) == p
     if ending_index is not None:
         assert isinstance(ending_index, list) and len(ending_index) == p
+
+    peak_mem = peak_mem or get_peak_mem(schedules)
+    max_len = check_and_get_schedule_len(schedules)
     starting_index = starting_index or [0] * p
     ending_index = ending_index or [max_len] * p
-
     last_index = [{_id: -1 for _id in "FfBbWw"} for _ in range(p)]
     for i in range(p):
         for j in range(max_len):
@@ -510,7 +554,6 @@ def reorder_greedily_without_increasing_peak_mem(schedules, m, starting_index =
                 continue
             last_index[i][identifier] = j
 
-    peak_mem = get_peak_mem(schedules)
     stage_mem = [0] * p
     def update_mem(stage_i, pass_c):
         if pass_c in "Ff":
@@ -645,6 +688,7 @@ def check_correctness(schedules, m, raise_exception=False):
                     return False
     return True
 
+
 def relabel_w(schedules, m):
     p = len(schedules)
     c_cnt = [{_id: 0 for _id in "FfBbWw"} for _ in range(p)]
@@ -654,7 +698,7 @@ def relabel_w(schedules, m):
                 continue
             c_cnt[i][schedules[i][j]] += 1
         for c in "FfBbWw":
-            assert c_cnt[i][c] == m
+            assert c_cnt[i][c] == m, f"{i}, {c}, {c_cnt[i][c]}"
     for i in range(p):
         w_queue = deque(maxlen=2 * m)
         for j in range(len(schedules[i])):
@@ -722,6 +766,8 @@ def schedule_by_building_block(p, m, building_block, max_mem, keep_stable_phase=
     if m < redundant_m:
         # 4. remove redundancy
         schedules = remove_redundancy(schedules, m)
+        if m <= p and 2 * m <= max_mem:
+            schedules = release_w_in_warmup_if_under_memory(schedules, peak_mem=min(2 * p, peak_mem))
         schedules = squeeze_without_change_order(schedules, m)
         print_schedules(schedules, "after removing redundancy")
         init_peak_mem = peak_mem = get_peak_mem(schedules)
@@ -820,6 +866,7 @@ def schedule(p, m, cost, max_mem):
         [4, -1, 4, -1],
         [5, -1, 5, -1]
     ]
+    # available_starting_patterns = available_starting_patterns[:1]
 
     best_schedule = None
     best_bubble = None