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