A better version

adaptive_schedule.py

@@ -1,5 +1,5 @@
 pattern_size = 6
-from collections import Counter
+from collections import Counter, deque
 from dataclasses import dataclass
 
 @dataclass(eq=True, frozen=True)
@@ -74,9 +74,6 @@ def transform_schedule(schedule, f, b, w, c):
     return result
 
 
-    
-
-
 def evaluate_schedule(schedule, f, b, w, c):
     stage_order = []
     local_prev = {}
@@ -123,7 +120,21 @@ def evaluate_schedule(schedule, f, b, w, c):
         r = max(get_time(sid, 'w', nmb - 1) - get_time(sid, 'F', 0) + f, r)
     return r
 
-def get_pattern_str(pos):
+
+debug = False
+def print_schedules(schedules, msg = None, force=False):
+    if not debug and not force:
+        return
+    if msg is not None:
+        print(msg)
+    for seq in schedules:
+        _str = ""
+        for v in seq:
+            _str += v
+        print(_str)
+
+
+def get_building_block_str(pos):
     pattern = [" "] * pattern_size
     notations = "FfBbWw"
     for i, v in enumerate(pos):
@@ -167,11 +178,11 @@ def calc_bubble(schedules):
     return stage_bubbles
 
 
-def init_repeated_schedule(p, m, patterns):
+def init_repeated_schedule(p, m, building_block):
     repeated = []
     _len = 4 * p + m + 1
     for i in range(p):
-        str_i = get_pattern_str(patterns[i]) * _len
+        str_i = get_building_block_str(building_block[i]) * _len
         repeated_i = []
         for v in str_i:
             repeated_i.append(v)
@@ -261,6 +272,8 @@ def process_warmup_without_increasing_peak_mem(schedules, m):
             elif char == 'W':
                 c_w += 1
             elif char == 'b':
+                break
+                # This logic can be removed because it is too complicated and should not impact the optimal solution
                 bj = j
                 while j < len(schedules[i]):
                     char = schedules[i][j]
@@ -290,8 +303,8 @@ def process_warmup_without_increasing_peak_mem(schedules, m):
             else:
                 assert char == ' '
             schedules[i][j] = ' '
-        assert c_f >= cnt_f[i] and c_ff >= cnt_ff[i]
-        assert c_w >= cnt_ff[p - 1] - cnt_ff[i] and c_b >= cnt_ff[p - 1] - cnt_ff[i]
+        # assert c_f >= cnt_f[i] and c_ff >= cnt_ff[i]
+        # assert c_w >= cnt_ff[p - 1] - cnt_ff[i] and c_b >= cnt_ff[p - 1] - cnt_ff[i]
         j = i
         u_f, u_ff, u_b, u_w = 0, 0, 0, 0
         for _ in range(2 * (p - 1 - i)):
@@ -365,15 +378,15 @@ def squeeze_without_change_order(schedules, m):
                         assert identifier_index[_cnt * p + i]['B'] >= 0
                     index = stage_index[i]
                 elif identifier in "FB":
-                    assert identifier_index[_cnt * p + i - 1][identifier] >= 0, "{} {} {}".format(i,identifier,_cnt)
+                    assert identifier_index[_cnt * p + i - 1][identifier] >= 0, "{} {} {}".format(i, identifier,_cnt)
                     index = max(identifier_index[_cnt * p + i - 1][identifier] + 1, stage_index[i])
                 elif identifier in "fb":
-                    assert identifier_index[_cnt * p + i + 1][identifier] >= 0, "{} {} {}".format(i,identifier,_cnt)
+                    assert identifier_index[_cnt * p + i + 1][identifier] >= 0, "{} {} {}".format(i, identifier,_cnt)
                     index = max(identifier_index[_cnt * p + i + 1][identifier] + 1, stage_index[i])
                 else:
                     raise
                 squeezed[i][index] = identifier
-                identifier_cnt[i][identifier] += 1
+                identifier_cnt[i][identifier] = _cnt + 1
                 identifier_index[_cnt * p + i][identifier] = index
                 stage_index[i] = index + 1
     return squeezed
@@ -396,7 +409,7 @@ def process_cooldown(schedules, m):
     p = len(schedules)
 
     peak_mem = get_peak_mem(schedules)
-    assert peak_mem <= 2 * p
+    assert peak_mem <= 2 * p, peak_mem
     max_bb = (peak_mem + 1) // 2
     max_bb = min(max_bb, m)
     max_b = min(peak_mem - max_bb, m)
@@ -406,7 +419,7 @@ def process_cooldown(schedules, m):
     for i in range(p):
         c_b, c_bb, c_w, c_ww = 0, 0, 0, 0
         last_ff_index = -1
-        # collect B/b which can be reorganized
+        # collect B/b which can be reordered
         for j in range(len(schedules[i]) - 1, -1, -1):
             char = schedules[i][j]
             if char == 'f' and last_ff_index == -1:
@@ -417,13 +430,15 @@ def process_cooldown(schedules, m):
             if char == 'b' and c_bb < max_bb:
                 schedules[i][j] = ' '
                 c_bb += 1
-        # clear W in the tail (#W + #w = peak_mem)
+        # clear W in the tail (#W + #w >= peak_mem & #W >= #B & #w >= #b)
         for j in range(len(schedules[i]) - 1, -1, -1):
             char = schedules[i][j]
-            if char == 'W' and c_w + c_ww < peak_mem:
+            if c_w >= c_b and c_ww >= c_bb and c_w + c_ww >= peak_mem:
+                break
+            if char == 'W':
                 schedules[i][j] = ' '
                 c_w += 1
-            if char == 'w' and c_w + c_ww < peak_mem:
+            if char == 'w':
                 schedules[i][j] = ' '
                 c_ww += 1
         if i == 0:
@@ -435,24 +450,17 @@ def process_cooldown(schedules, m):
             schedules[i][index] = 'b'
         for k in range(c_b):
             index = starting_index + 1 + i - 2 * k
-            assert schedules[i][index] == ' ', schedules[i][index]
+            # assert schedules[i][index] == ' ', schedules[i][index]
             schedules[i][index] = 'B'
 
-    # 2: squeeze cooldown phase without change order
-    schedules = squeeze_without_change_order(schedules, m)
-
-    # 3: add W back in cooldown phase
+    # 2: add W back in cooldown phase
     for i in range(p):
         c_w, c_ww = 0, 0
-        last_w_index = -2
+        last_w_index = -1
         for j in range(len(schedules[i]) - 1, -1, -1):
             if schedules[i][j] in "Ww":
-                if last_w_index < 0:
-                    schedules[i][j] = ' '
-                    last_w_index += 1
-                else:
-                    last_w_index = j
-                    break
+                last_w_index = j
+                break
         for j in range(len(schedules[i])):
             char = schedules[i][j]
             if char == 'B':
@@ -475,38 +483,281 @@ def process_cooldown(schedules, m):
     return schedules
 
 
-def schedule_by_pattern(p, m, patterns, max_mem):
-    schedules = init_repeated_schedule(p, max(m, 2 * p), patterns)
-    schedules = clear_invalid_index(schedules,  max(m, 2 * p))
-    init_peak_mem = get_peak_mem(schedules)
-    if init_peak_mem > max_mem:
-        return None, init_peak_mem, [6 *  max(m, 2 * p)] * p
-    schedules = process_warmup_without_increasing_peak_mem(schedules,  max(m, 2 * p))
-    
+def reorder_greedily_without_increasing_peak_mem(schedules, m, starting_index = None, ending_index = 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,
+    and check whether it is possible to move if we keep all other passes unchanged. If the check succeeds, we move it
+    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
+    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):
+            identifier = schedules[i][j]
+            if identifier == ' ':
+                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":
+            stage_mem[stage_i] += 1
+        elif pass_c in "Ww":
+            stage_mem[stage_i] -= 1
+
+    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)]
+    for j in range(0, max_len):
+        for i in range(p):
+            identifier = schedules[i][j]
+            if identifier in "FfBbWw":
+                _cnt = identifier_cnt[i][identifier]
+                identifier_cnt[i][identifier] = _cnt + 1
+                identifier_index[_cnt * p + i][identifier] = j
+                update_mem(i, identifier)
+                continue
+            assert identifier == ' '
+            if j < starting_index[i] or j >= ending_index[i]:
+                continue
+            available = set()
+            for c in "FfBbWw":
+                if last_index[i][c] > j:
+                    available.add(c)
+            mem_delta, peak_delta = 0, 0
+            for k in range(j + 1, ending_index[i]):
+                if len(available) == 0:
+                    break
+                identifier = schedules[i][k]
+                if identifier in "Ff":
+                    mem_delta += 1
+                elif identifier in "Ww":
+                    mem_delta -= 1
+                prev_peak = peak_delta
+                peak_delta = max(peak_delta, mem_delta)
+                if identifier == ' ' or identifier not in available:
+                    continue
+                available.remove(identifier)
+                if identifier in "Ff" and stage_mem[i] + prev_peak >= peak_mem:
+                    # will increase peak memory
+                    continue
+                can_move = True
+                _cnt = identifier_cnt[i][identifier]
+                if identifier in "FB":
+                    if i > 0:
+                        _index = identifier_index[_cnt * p + i - 1][identifier]
+                        if _index <= -1 or _index >= j:
+                            can_move = False
+                    elif identifier == 'B':
+                        if identifier_cnt[i]['f'] <= _cnt:
+                            can_move = False
+                elif identifier in "fb":
+                    if i + 1 < p:
+                        _index = identifier_index[_cnt * p + i + 1][identifier]
+                        if _index <= -1 or _index >= j:
+                            can_move = False
+                    else:
+                        _pi = 'F' if identifier == 'f' else 'B'
+                        if identifier_cnt[i][_pi] <= _cnt:
+                            can_move = False
+                elif identifier in "Ww":
+                    _bi = 'B' if identifier == 'W' else 'b'
+                    if identifier_cnt[i][_bi] <= _cnt:
+                        can_move = False
+                else:
+                    assert False
+                if not can_move:
+                    continue
+                # if i == 0:
+                #     print(peak_mem, stage_mem[i], identifier, mem_delta)
+                schedules[i][j] = identifier
+                schedules[i][k] = ' '
+                identifier_cnt[i][identifier] = _cnt + 1
+                identifier_index[_cnt * p + i][identifier] = j
+                update_mem(i, identifier)
+                break
+    return schedules
+
+
+def check_correctness(schedules, m, raise_exception=False):
+    p = len(schedules)
+    c_index = [{_id: -1 for _id in "FfBbWw"} for _ in range(p * m)]
+    for i in range(p):
+        c_cnt = {_id: 0 for _id in "FfBbWw"}
+        for j in range(len(schedules[i])):
+            c = schedules[i][j]
+            if c in "FfBbWw":
+                _cnt = c_cnt[c]
+                assert _cnt < m
+                c_index[_cnt * p + i][c] = j
+                c_cnt[c] = _cnt + 1
+        for c in "FfBbWw":
+            if c_cnt[c] != m:
+                assert not raise_exception
+                return False
+    for i in range(p):
+        for j in range(m):
+            for c in "FfBbWw":
+                if c_index[j * p + i][c] == -1:
+                    assert not raise_exception
+                    return False
+            if c_index[j * p + i]['B'] >= c_index[j * p + i]['W']:
+                assert not raise_exception, f"{i} {j} {c}"
+                return False
+            if c_index[j * p + i]['b'] >= c_index[j * p + i]['w']:
+                assert not raise_exception
+                return False
+            if i == 0:
+                if c_index[j * p + i]['f'] >= c_index[j * p + i]['B']:
+                    assert not raise_exception
+                    return False
+            elif i == p - 1:
+                if c_index[j * p + i]['F'] >= c_index[j * p + i]['f']:
+                    assert not raise_exception
+                    return False
+                if c_index[j * p + i]['B'] >= c_index[j * p + i]['b']:
+                    assert not raise_exception
+                    return False
+            else:
+                if c_index[j * p + i - 1]['F'] >= c_index[j * p + i]['F']:
+                    assert not raise_exception
+                    return False
+                if c_index[j * p + i - 1]['B'] >= c_index[j * p + i]['B']:
+                    assert not raise_exception
+                    return False
+                if c_index[j * p + i + 1]['f'] >= c_index[j * p + i]['f']:
+                    assert not raise_exception
+                    return False
+                if c_index[j * p + i + 1]['b'] >= c_index[j * p + i]['b']:
+                    assert not raise_exception
+                    return False
+    return True
+
+def relabel_w(schedules, m):
+    p = len(schedules)
+    c_cnt = [{_id: 0 for _id in "FfBbWw"} for _ in range(p)]
+    for i in range(p):
+        for j in range(len(schedules[i])):
+            if schedules[i][j] == ' ':
+                continue
+            c_cnt[i][schedules[i][j]] += 1
+        for c in "FfBbWw":
+            assert c_cnt[i][c] == m
+    for i in range(p):
+        w_queue = deque(maxlen=2 * m)
+        for j in range(len(schedules[i])):
+            identifier = schedules[i][j]
+            if identifier == 'B':
+                w_queue.append('W')
+            elif identifier == 'b':
+                w_queue.append('w')
+            elif identifier in "Ww":
+                assert len(w_queue) > 0, f"{i} {j}"
+                schedules[i][j] = w_queue.popleft()
+        assert len(w_queue) == 0
+    return schedules
+
+
+def remove_redundancy(schedules, m):
     for sid in range(len(schedules)):
         cnt = {_id: 0 for _id in "FfBbWw"}
         for i in range(len(schedules[sid])):
-            if(schedules[sid][i] == ' '):
+            if schedules[sid][i] == ' ':
                 continue
             if cnt[schedules[sid][i]] >= m:
                 schedules[sid][i] = ' '
             else:
                 cnt[schedules[sid][i]] += 1
+    return schedules
+
+
+def schedule_by_building_block(p, m, building_block, max_mem, keep_stable_phase=False):
+    # Apply the framework of repeating-squeezing-reordering
+    # 1. repeating
+    redundant_m = max(m, 2 * p)  # we add some redundant micro-batches to avoid unexpected bugs
+    schedules = init_repeated_schedule(p, redundant_m, building_block)
+    schedules = clear_invalid_index(schedules,  redundant_m)
+    init_peak_mem = get_peak_mem(schedules)
+    if (m == redundant_m and init_peak_mem > max_mem) or init_peak_mem > 2 * p:
+        return None, init_peak_mem, [6 * m] * p
+    print_schedules(schedules, "after repeating")
+
+    # 2. squeezing
+    schedules = squeeze_without_change_order(schedules, redundant_m)
+    print_schedules(schedules, "after squeezing")
+
+    # 3. reordering
+    # 3.a. reorder warm-up
+    schedules = process_warmup_without_increasing_peak_mem(schedules,  redundant_m)  # must work with m >= 2p
+    schedules = squeeze_without_change_order(schedules, redundant_m)
+    if keep_stable_phase:
+        ending_index = [0] * p  # before second b
+        for i in range(p):
+            bb_cnt = 0
+            for j in range(len(schedules[i])):
+                if schedules[i][j] == 'b':
+                    bb_cnt += 1
+                    if bb_cnt >= 2:
+                        ending_index[i] = j
+                        break
+        schedules = reorder_greedily_without_increasing_peak_mem(schedules, redundant_m, ending_index=ending_index)
     peak_mem = get_peak_mem(schedules)
+    if debug:
+        assert peak_mem <= init_peak_mem, f"{init_peak_mem}, {peak_mem}"
     if peak_mem > init_peak_mem:
         return None, init_peak_mem, [6 * m] * p
 
-    schedules = squeeze_without_change_order(schedules, m)
+    if m < redundant_m:
+        # 4. remove redundancy
+        schedules = remove_redundancy(schedules, m)
+        schedules = squeeze_without_change_order(schedules, m)
+        print_schedules(schedules, "after removing redundancy")
+        init_peak_mem = peak_mem = get_peak_mem(schedules)
+        if peak_mem > max_mem:
+            return None, peak_mem, [6 * m] * p
 
+    # 3.b. reorder cool-down
     schedules = process_cooldown(schedules, m)
+    if keep_stable_phase:
+        starting_index = [0] * p
+        for i in range(p):
+            for j in range(len(schedules[i])):
+                if schedules[i][j] == 'F':
+                    starting_index[i] = j
+        schedules = reorder_greedily_without_increasing_peak_mem(schedules, m, starting_index=starting_index)
+    if not keep_stable_phase:
+        reorder_greedily_without_increasing_peak_mem(schedules, m)
+    schedules = relabel_w(schedules, m)
+    print_schedules(schedules, "after reordering")
     peak_mem = get_peak_mem(schedules)
+    if debug:
+        assert peak_mem <= init_peak_mem, f"{init_peak_mem}, {peak_mem}"
     if peak_mem > init_peak_mem:
         return None, init_peak_mem, [6 * m] * p
+
+    # return
+    if not check_correctness(schedules, m, raise_exception=debug):
+        return None, peak_mem, [6 * m] * p
     stage_bubbles = calc_bubble(schedules)
+    if debug:
+        print(peak_mem, stage_bubbles)
+        print("-" * 100)
     return schedules, peak_mem, stage_bubbles
 
 
-def fill_w_in_pattern(pattern):
+def fill_w_in_building_block(pattern):
     f, ff, b, bb, w, ww = 0, 1, 2, 3, 4, 5
     vis = [False] * pattern_size
     for v in pattern:
@@ -523,10 +774,11 @@ def fill_w_in_pattern(pattern):
     return pattern
 
 
-def get_whole_pattern(pattern_0, offset_0, offset_1, len_0, p):
-    whole_pattern = [pattern_0]
+def get_building_block(pattern_0, offset_0, offset_1, len_0, p):
+    # see Appendix A in the paper
+    build_block = [pattern_0]
     for i in range(p - 1):
-        last_pattern = whole_pattern[i]
+        last_pattern = build_block[i]
         new_pattern = [-1] * pattern_size
         vis = [False] * pattern_size
         if i < len_0:
@@ -540,26 +792,28 @@ def get_whole_pattern(pattern_0, offset_0, offset_1, len_0, p):
                 return None
             vis[pos] = True
             new_pattern[v] = pos
-        new_pattern = fill_w_in_pattern(new_pattern)
-        whole_pattern.append(new_pattern)
-    return whole_pattern
+        new_pattern = fill_w_in_building_block(new_pattern)
+        build_block.append(new_pattern)
+    return build_block
 
 
         
 def schedule(p, m, cost, max_mem):
     f, ff, b, bb, w, ww = 0, 1, 2, 3, 4, 5
-    available_patterns = []
+    available_starting_patterns = []
+    # iterate available patterns for the first row/device of a building block
     for ff_i in range(1, pattern_size):
         for b_i in range(1, pattern_size):
             for bb_i in range(1, pattern_size):
                 if ff_i == b_i or ff_i == bb_i or b_i == bb_i:
                     continue
                 pattern = [0, ff_i, b_i, bb_i, -1, -1]
-                pattern = fill_w_in_pattern(pattern)
-                available_patterns.append(pattern)
+                pattern = fill_w_in_building_block(pattern)
+                available_starting_patterns.append(pattern)
 
-    print(len(available_patterns))
+    # available uniform offsets, see Section 3.1 in the paper.
     available_offsets = [
+        # [\delta_F^0, \delta_F^1, \delta_B^1, \delta_B^0]
         [1, -1, 1, -1],
         [2, -1, 2, -1],
         [3, -1, 3, -1],
@@ -569,23 +823,37 @@ def schedule(p, m, cost, max_mem):
 
     best_schedule = None
     best_bubble = None
-    for pattern_0 in available_patterns:
+    peak_mem2min_bubble = {}
+    for pattern_0 in available_starting_patterns:
         for i_0 in range(len(available_offsets)):
             for i_1 in range(i_0 + 1):
                 for len_0 in range(1, p):
                     offset_0 = available_offsets[i_0]
                     offset_1 = available_offsets[i_1]
-                    whole_pattern = get_whole_pattern(pattern_0, offset_0, offset_1, len_0, p)
-                    if whole_pattern is None:
+                    build_block = get_building_block(pattern_0, offset_0, offset_1, len_0, p)
+                    if build_block is None:
                         continue
-                    s, peak_mem, bubbles = schedule_by_pattern(p, m, whole_pattern, min(2 * p, max_mem))
+                    s, peak_mem, bubbles = schedule_by_building_block(p, m, build_block, min(2 * p, max_mem))
                     if peak_mem > 2 * p or peak_mem > max_mem:
                         break
                     if s is None:
                         continue
-                    max_bubble = max(bubbles)
                     max_bubble = evaluate_schedule(s, *cost)
                     if best_schedule is None or max_bubble < best_bubble:
                         best_schedule, best_bubble = s, max_bubble
+
+                    max_bubble = max(bubbles)
+                    min_bubble = min(peak_mem2min_bubble.get(peak_mem, max_bubble), max_bubble)
+                    peak_mem2min_bubble[peak_mem] = min_bubble
+    mem2bubble = {}
+    for peak_mem in sorted(peak_mem2min_bubble.keys()):
+        bubble = peak_mem2min_bubble[peak_mem]
+        mem2bubble[peak_mem] = bubble
+        # expected_bubble = max(0, 6 * p - 1 - 3 * peak_mem)
+        expected_bubble = 3 * p - 1 - 3 * peak_mem + max(3 * p, p - 1 + (1+(peak_mem+1)//2)*2)
+        # expected_bubble = 6 * p - 1 - 3 * peak_mem
+        print(peak_mem, bubble, expected_bubble, "|", bubble - expected_bubble)
+    print(mem2bubble)
+
     res = transform_schedule(best_schedule, *cost)
-    return res
+    return res