def __call__(self, model, n_chips):
board = Board()
board.new_chip()
def get_chip(i):
chip = board.chips[-1]
assert len(chip.cores) <= self.cores_per_chip
if len(chip.cores) == self.cores_per_chip:
assert len(board.chips) < n_chips, (
"The network needs more chips than requested (%d)" % n_chips,
)
chip = board.new_chip()
return chip
i = 0
for input in model.inputs:
self.input_to_board(input, board)
i += 1
for block in model.blocks:
self.block_to_new_core(block, get_chip(i))
i += 1
board.probes.extend(model.probes)
logger.info("Greedy allocation across %d chips", board.n_chips)
return board
def __call__(self, model, n_chips):
block_map = dict(enumerate(model.blocks))
block_rates = None
if self.ensemble_rates is not None:
block_rates = ensemble_to_block_rates(model, self.ensemble_rates)
block_rates = {
block: np.round(rate * self.rate_scale)
for block, rate in block_rates.items()
}
block_conns = estimate_interblock_activity(block_map, block_rates=block_rates)
# partition graph
G = networkx.Graph()
G.add_nodes_from(block_map.keys())
edge_map = set()
for i in block_map:
for j, val in block_conns[i].items():
if (i, j) in edge_map or (j, i) in edge_map:
continue
val = val + block_conns[j].get(i, 0)
G.add_edge(i, j, weight=int(round(val))) # weights must be integers
edge_map.add((i, j))
edge_map.add((j, i))
_, parts = nxmetis.partition(G, nparts=int(n_chips))
for i, part in enumerate(parts):
if len(part) > 128:
raise ValueError(
f"Partition {i} has {len(part)} cores, "
"which exceeds the available 128 cores"
)
# --- create board
board = Board()
# add inputs to board
for input in model.inputs:
self.input_to_board(input, board)
# blocks to chips
for part in parts:
chip = board.new_chip()
for block_idx in part:
block = block_map[block_idx]
self.block_to_new_core(block, chip)
# add probes
board.probes.extend(model.probes)
logger.info("METIS allocation across %d chips", board.n_chips)
return board
def __call__(self, model):
board = Board()
chip = board.new_chip()
for block in model.blocks:
self.block_to_chip(block, chip)
for input in model.inputs:
self.input_to_chip(input, chip)
return board
def __call__(self, model, n_chips): # noqa: C901
block_map = dict(enumerate(model.blocks))
block_rates = (
ensemble_to_block_rates(model, self.ensemble_rates)
if self.ensemble_rates is not None
else None
)
block_conns_out = estimate_interblock_activity(
block_map, block_rates=block_rates
)
block_conns_in = {
i: {j: block_conns_out[j][i] for j in block_map if i in block_conns_out[j]}
for i in block_map
}
# find blocks with no pre block
no_pre_blocks = []
for i in block_map:
if sum(v for v in block_conns_in[i].values()) == 0:
no_pre_blocks.append(i)
# --- create board
board = Board()
# add inputs to board
for input in model.inputs:
self.input_to_board(input, board)
# --- add blocks to chips
chip = None
unallocated_blocks = set(block_map)
while len(unallocated_blocks) > 0:
if chip is None or len(chip.cores) == self.cores_per_chip:
assert (
len(board.chips) < n_chips
), f"The network needs more chips than requested ({n_chips})"
# start a new chip
chip = board.new_chip()
# choose a no-pre block, if possible
for block_idx in no_pre_blocks:
if block_idx in unallocated_blocks:
break
else:
block_idx = next(iter(unallocated_blocks))
chip_blocks = set()
else:
# choose the block with the largest connection to blocks on this chip
block_idx = -1
max_conn = 0
for i in chip_blocks:
for j in unallocated_blocks.intersection(block_conns_out[i]):
ij = block_conns_out[i][j]
if ij > max_conn:
max_conn = ij
block_idx = j
for j in unallocated_blocks.intersection(block_conns_in[i]):
ij = block_conns_in[i][j]
if ij > max_conn:
max_conn = ij
block_idx = j
if block_idx < 0:
# none of the remaining blocks connect to blocks on this chip,
# so pick a no-pre block if possible, otherwise any block will do.
for block_idx in no_pre_blocks:
if block_idx in unallocated_blocks:
break
else:
block_idx = next(iter(unallocated_blocks))
block = block_map[block_idx]
self.block_to_new_core(block, chip)
chip_blocks.add(block_idx)
unallocated_blocks.remove(block_idx)
# add probes
board.probes.extend(model.probes)
logger.info("GreedyInterchip allocation across %d chips", board.n_chips)
return board