pattern_size = 6
from collections import Counter, deque
from dataclasses import dataclass
@dataclass(eq=True, frozen=True)
class ScheduledNode:
type: str
chunk: int
stage: int
minibatch: int
start_time: int
completion_time: int
def transform_schedule(schedule, f, b, w, c):
result = []
stage_order = []
local_prev = {}
stages = len(schedule)
for sid, stage in enumerate(schedule):
counter = Counter()
order = []
for p in stage:
if not p.strip():
continue
mb = counter.get(p, 0)
if order:
local_prev[(sid, p, mb)] = order[-1]
order.append((p, mb))
counter.update(p)
stage_order.append(order)
nmb = max(counter.values())
time_map = {}
cost = {
'F': f,
'B': b,
'W': w,
'f': f,
'b': b,
'w': w,
}
def get_time(stage, type, mb):
if (stage, type, mb) in time_map:
return time_map.get((stage, type, mb))
time = 0
if (stage, type, mb) in local_prev:
time = get_time(stage, *local_prev[(stage, type, mb)])
if type in "FB" and stage > 0:
time = max(time, get_time(stage - 1, type, mb) + c)
if type in "fb" and stage + 1< len(schedule):
time = max(time, get_time(stage + 1, type, mb) + c)
# print(f'{stage} {type}:{mb}', time + cost[type])
time_map[(stage, type, mb)] = time + cost[type]
return time_map[(stage, type, mb)]
r = 0
for sid, stage in enumerate(schedule):
r = max(get_time(sid, 'W', nmb - 1) - get_time(sid, 'F', 0) + f, r)
r = max(get_time(sid, 'w', nmb - 1) - get_time(sid, 'F', 0) + f, r)
for sid, stage in enumerate(stage_order):
result_stage = []
for p, mb in stage:
result_stage.append(ScheduledNode(
p.upper(),
p in "fBW",
sid,
mb,
get_time(sid, p, mb) - cost[p],
get_time(sid, p, mb)
)
)
result.append(result_stage)
return result
def evaluate_schedule(schedule, f, b, w, c):
stage_order = []
local_prev = {}
stages = len(schedule)
for sid, stage in enumerate(schedule):
counter = Counter()
order = []
for p in stage:
if not p.strip():
continue
mb = counter.get(p, 0)
if order:
local_prev[(sid, p, mb)] = order[-1]
order.append((p, mb))
counter.update(p)
stage_order.append(order)
nmb = max(counter.values())
time_map = {}
cost = {
'F': f,
'B': b,
'W': w,
'f': f,
'b': b,
'w': w,
}
def get_time(stage, type, mb):
if (stage, type, mb) in time_map:
return time_map.get((stage, type, mb))
time = 0
if (stage, type, mb) in local_prev:
time = get_time(stage, *local_prev[(stage, type, mb)])
if type in "FB" and stage > 0:
time = max(time, get_time(stage - 1, type, mb) + c)
if type in "fb" and stage + 1< len(schedule):
time = max(time, get_time(stage + 1, type, mb) + c)
# print(f'{stage} {type}:{mb}', time + cost[type])
time_map[(stage, type, mb)] = time + cost[type]
return time_map[(stage, type, mb)]
r = 0
for sid, stage in enumerate(schedule):
r = max(get_time(sid, 'W', nmb - 1) - get_time(sid, 'F', 0) + f, r)
r = max(get_time(sid, 'w', nmb - 1) - get_time(sid, 'F', 0) + f, r)
return r
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):
if v < 0:
continue
pattern[v] = notations[i]
_str = ""
for v in pattern:
_str += v
return _str
def get_peak_mem(schedules, return_all=False):
max_peak = 0
all_peak = []
for schedule_ in schedules:
peak, mem = 0, 0
for v in schedule_:
if v in "Ff":
mem += 1
elif v in "Ww":
mem -= 1
peak = max(peak, mem)
all_peak.append(peak)
max_peak = max(max_peak, peak)
if return_all:
return all_peak
return max_peak
def calc_bubble(schedules):
stage_bubbles = []
for i in range(len(schedules)):
max_len = 0
count = 0
for j in range(len(schedules[i])):
if schedules[i][j] != ' ':
max_len = j + 1
count += 1
stage_bubbles.append(max_len - count - i)
return stage_bubbles
def init_repeated_schedule(p, m, building_block):
repeated = []
_len = 4 * p + m + 1
for i in range(p):
str_i = get_building_block_str(building_block[i]) * _len
repeated_i = []
for v in str_i:
repeated_i.append(v)
repeated.append(repeated_i)
return repeated
def clear_invalid(repeated, stage, pos, offset=-1):
while 0 <= pos < len(repeated[stage]):
repeated[stage][pos] = ' '
pos += offset * pattern_size
return repeated
def clear_invalid_index(repeated, m):
p = len(repeated)
index = pattern_size
for identifier in "FfBb":
if identifier in "FB":
_iter = range(p)
else:
_iter = range(p - 1, -1, -1)
for i in _iter:
for j in range(pattern_size):
if repeated[i][index] == identifier:
clear_invalid(repeated, i, index - pattern_size, offset=-1)
clear_invalid(repeated, i, index + pattern_size * m, offset=1)
index += 1
if identifier in "Bb":
w_identifier = {'B': 'W', 'b': 'w'}[identifier]
for k in range(pattern_size):
if repeated[i][index + k] == w_identifier:
clear_invalid(repeated, i, index + k - pattern_size, offset=-1)
clear_invalid(repeated, i, index + k + pattern_size * m, offset=1)
break
break
index += 1
return repeated
def process_warmup_without_increasing_peak_mem(schedules, m):
"""
FFFFFFFFFF fBWfBWfBWfBWfBW b
FFFFFFFFF f fBWfBWfBWfBWFBWb
FFFFFFFF f f fBWfBWfBWFBW b
FFFFFFF f f f fBWfBWFBW Bb
FFFFFF f f f f fBWFBWFBWb
FFFFFfFf f f f BWFBW b
FFFfFfFfFf f BW Bb
FfFfFfFfFfF BWb
We reorganize the warmup phase in the following way (i -> pipeline stage from 0):
1. Before the first B, we set #f = min(i+1, peak_mem//2), #F = peak_mem - #f
2. Before the first b, #f = peak_mem//2
3. The offset between the first B is 1
4. Before the first b, we use the pattern of (BWf)*j + (BWF)*k,
where j = max(0, peak_mem//2 - (i+1)), k = max(0, #W - j - 1)
"""
# process warmup phase (before the first b)
p = len(schedules)
peak_mem = get_peak_mem(schedules)
peak_mem = min(peak_mem, 2 * p)
cnt_f, cnt_ff = [], []
for i in range(p):
cc_ff = min(i + 1, peak_mem // 2)
cc_ff = min(cc_ff, m)
cc_f = min(peak_mem - cc_ff, m)
cnt_f.append(cc_f)
cnt_ff.append(cc_ff)
distance_b2bb = 0
for j in range(len(schedules[p - 1])):
if schedules[p - 1][j] == 'B':
for k in range(j, len(schedules[p - 1])):
if schedules[p - 1][k] == 'b':
distance_b2bb = k - j
break
break
for i in range(p):
c_f, c_ff, c_b, c_w = 0, 0, 0, 0
for j in range(len(schedules[i])):
char = schedules[i][j]
if char == 'F':
c_f += 1
elif char == 'f':
c_ff += 1
elif char == 'B':
c_b += 1
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]
if char == 'f' and c_ff < cnt_ff[p - 1]:
schedules[i][j] = ' '
c_ff += 1
if char == 'B' and c_b < c_ff:
if c_b < (2 * (p - i) + distance_b2bb) // 3 or c_b < cnt_ff[p - 1] - cnt_ff[i]:
# there is empty space, or the number of B is not enough to cover extra f
schedules[i][j] = ' '
c_b += 1
if char == 'W' and c_w < c_b:
if c_w < (2 * (p - i) + distance_b2bb - 1) // 3 or c_w < cnt_ff[p - 1] - cnt_ff[i]:
# there is empty space, or the number of W is not enough to cover extra f
schedules[i][j] = ' '
c_w += 1
j += 1
j = bj
while j < len(schedules[i]):
if schedules[i][j] == 'F':
if c_f < c_ff or c_f < cnt_f[i] or c_f - cnt_f[i] + c_ff - cnt_ff[i] < c_w - 1:
# put enough F, or there are some unused BW
schedules[i][j] = ' '
c_f += 1
j += 1
break
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]
j = i
u_f, u_ff, u_b, u_w = 0, 0, 0, 0
for _ in range(2 * (p - 1 - i)):
if u_f < cnt_f[i] and u_f < c_f:
schedules[i][j] = 'F'
u_f += 1
j += 1
for _ in range(i + 1):
if u_f < cnt_f[i] and u_f < c_f:
schedules[i][j] = 'F'
u_f += 1
j += 1
if u_ff < cnt_ff[i] and u_ff < c_ff:
schedules[i][j] = 'f'
u_ff += 1
j += 1
while u_f < c_f or u_ff < c_ff or u_b < c_b or u_w < c_w:
if u_b < c_b:
schedules[i][j] = 'B'
u_b += 1
j += 1
if u_w < c_w:
schedules[i][j] = 'W'
u_w += 1
j += 1
if u_ff < c_ff:
assert u_ff < u_f
schedules[i][j] = 'f'
u_ff += 1
elif u_f < c_f:
schedules[i][j] = 'F'
u_f += 1
j += 1
return schedules
def squeeze_without_change_order(schedules, m):
p = len(schedules)
squeezed = [[' '] * len(schedules[_]) for _ in range(p)]
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)]
stage_index = [0 for _ in range(p)]
for j in range(max_len):
for _dir in range(2):
if _dir == 0:
_iter = range(p)
else:
_iter = range(p - 1, -1, -1)
for i in _iter:
identifier = schedules[i][j]
if identifier == ' ':
continue
if _dir == 0 and identifier in "fbw":
continue
if _dir == 1 and identifier in "FBW":
continue
_cnt = identifier_cnt[i][identifier]
assert _cnt < m, "{} - {}, {}".format(i, identifier, _cnt)
if identifier in "Ww" or (i == 0 and identifier in "FB") or (i == p - 1 and identifier in "fb"):
if i == 0 and identifier == 'B':
assert identifier_index[_cnt * p + i]['f'] >= 0
if i == p - 1 and identifier == 'f':
assert identifier_index[_cnt * p + i]['F'] >= 0
if i == p - 1 and identifier == 'b':
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)
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)
index = max(identifier_index[_cnt * p + i + 1][identifier] + 1, stage_index[i])
else:
raise
squeezed[i][index] = identifier
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
def process_cooldown(schedules, m):
"""
fBW bwbwbwbw
fBWBW bwbwbwbw
fBWBWBW bwbwbwbw
fBWBWBWBW bwbwbwbw
f BWBWBWBbWbwbwbww
f BWBWBbBbWbWbwwww
f BWBbBbBbWbWWwwww
f BbBbBbBbWWWWwwww
We reorganize the cooldown phase in the following way (i -> pipeline stage from 0):
1. After the last f, we set #b = (peak_mem+1)//2, and #B = min(i+1, peak_mem - #b)
2. After the last f, we make all the dependencies as tight as possible
"""
p = len(schedules)
peak_mem = get_peak_mem(schedules)
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)
# 1: reorganize B/b and remove W/w in cooldown phase
starting_index = -1
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 reordered
for j in range(len(schedules[i]) - 1, -1, -1):
char = schedules[i][j]
if char == 'f' and last_ff_index == -1:
last_ff_index = j
if char == 'B' and c_b < i + 1 and c_b < max_b:
schedules[i][j] = ' '
c_b += 1
if char == 'b' and c_bb < max_bb:
schedules[i][j] = ' '
c_bb += 1
# 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 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':
schedules[i][j] = ' '
c_ww += 1
if i == 0:
starting_index = last_ff_index
# reorganize B/b in the tail
for k in range(c_bb):
index = starting_index - i + 2 * p - 2 * k
assert schedules[i][index] == ' ', "{} {} {}".format(schedules[i][index], k, i)
schedules[i][index] = 'b'
for k in range(c_b):
index = starting_index + 1 + i - 2 * k
# assert schedules[i][index] == ' ', schedules[i][index]
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
for j in range(len(schedules[i]) - 1, -1, -1):
if schedules[i][j] in "Ww":
last_w_index = j
break
for j in range(len(schedules[i])):
char = schedules[i][j]
if char == 'B':
c_w += 1
elif char == 'b':
c_ww += 1
elif char == 'W':
c_w -= 1
elif char == 'w':
c_ww -= 1
if char == ' ' and j > last_w_index:
if c_w > 0:
schedules[i][j] = 'W'
c_w -= 1
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 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,
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)
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):
identifier = schedules[i][j]
if identifier == ' ':
continue
last_index[i][identifier] = j
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, 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])):
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] == ' ':
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
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)
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_building_block(pattern):
f, ff, b, bb, w, ww = 0, 1, 2, 3, 4, 5
vis = [False] * pattern_size
for v in pattern:
if v >= 0:
vis[v] = True
assert pattern[b] >= 0 and pattern[bb] >= 0
for v, vw in [(b, w), (bb, ww)]:
for j in range(pattern_size):
pos = (pattern[v] + j) % pattern_size
if not vis[pos]:
pattern[vw] = pos
vis[pos] = True
break
return pattern
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 = build_block[i]
new_pattern = [-1] * pattern_size
vis = [False] * pattern_size
if i < len_0:
offset = offset_0
else:
offset = offset_1
for v, v_o in enumerate(offset):
pos = (last_pattern[v] + v_o + pattern_size) % pattern_size
assert 0 <= pos < pattern_size
if vis[pos]:
return None
vis[pos] = True
new_pattern[v] = pos
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_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_building_block(pattern)
available_starting_patterns.append(pattern)
# 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],
[4, -1, 4, -1],
[5, -1, 5, -1]
]
# available_starting_patterns = available_starting_patterns[:1]
best_schedule = None
best_bubble = None
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]
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_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 = 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