def _setup_transfers(group): """ Compute all transfers that are owned by our parent group. Parameters ---------- group : <Group> Parent group. """ rev = group._mode != 'fwd' for subsys in group._subgroups_myproc: subsys._setup_transfers() abs2meta = group._var_abs2meta allprocs_abs2meta = group._var_allprocs_abs2meta myproc = group.comm.rank transfers = group._transfers = {} vectors = group._vectors offsets = group._get_var_offsets() vec_names = group._lin_rel_vec_name_list if group._use_derivatives else group._vec_names mypathlen = len(group.pathname + '.' if group.pathname else '') sub_inds = group._subsystems_inds for vec_name in vec_names: relvars, _ = group._relevant[vec_name]['@all'] # Initialize empty lists for the transfer indices nsub_allprocs = len(group._subsystems_allprocs) xfer_in = [] xfer_out = [] fwd_xfer_in = [[] for s in group._subsystems_allprocs] fwd_xfer_out = [[] for s in group._subsystems_allprocs] if rev: rev_xfer_in = [[] for s in group._subsystems_allprocs] rev_xfer_out = [[] for s in group._subsystems_allprocs] allprocs_abs2idx = group._var_allprocs_abs2idx[vec_name] sizes_in = group._var_sizes[vec_name]['input'] sizes_out = group._var_sizes[vec_name]['output'] offsets_in = offsets[vec_name]['input'] offsets_out = offsets[vec_name]['output'] # Loop through all connections owned by this system for abs_in, abs_out in group._conn_abs_in2out.items(): if abs_out not in relvars['output'] or abs_in not in relvars[ 'input']: continue # Only continue if the input exists on this processor if abs_in in abs2meta: # Get meta meta_in = abs2meta[abs_in] meta_out = allprocs_abs2meta[abs_out] idx_in = allprocs_abs2idx[abs_in] idx_out = allprocs_abs2idx[abs_out] # Read in and process src_indices src_indices = meta_in['src_indices'] if src_indices is None: owner = group._owning_rank[abs_out] # if the input is larger than the output on a single proc, we have # to just loop over the procs in the same way we do when src_indices # is defined. if meta_in['size'] > sizes_out[owner, idx_out]: src_indices = np.arange(meta_in['size'], dtype=INT_DTYPE) elif src_indices.ndim == 1: if _is_slice(src_indices): indices = _slice_indices(src_indices, meta_out['global_size'], meta_out['global_shape']) src_indices = convert_neg(indices, meta_out['global_size']) else: src_indices = _flatten_src_indices( src_indices, meta_in['shape'], meta_out['global_shape'], meta_out['global_size']) # 1. Compute the output indices # NOTE: src_indices are relative to a single, possibly distributed variable, # while the output_inds that we compute are relative to the full distributed # array that contains all local variables from each rank stacked in rank order. if src_indices is None: if meta_out['distributed']: # input in this case is non-distributed (else src_indices would be # defined by now). The input size must match the full # distributed size of the output. for rank, sz in enumerate(sizes_out[:, idx_out]): if sz > 0: out_offset = offsets_out[rank, idx_out] break output_inds = np.arange(out_offset, out_offset + meta_out['global_size'], dtype=INT_DTYPE) else: rank = myproc if abs_out in abs2meta else owner offset = offsets_out[rank, idx_out] output_inds = np.arange(offset, offset + meta_in['size'], dtype=INT_DTYPE) else: output_inds = np.zeros(src_indices.size, INT_DTYPE) start = end = 0 for iproc in range(group.comm.size): end += sizes_out[iproc, idx_out] if start == end: continue # The part of src on iproc on_iproc = np.logical_and(start <= src_indices, src_indices < end) if np.any(on_iproc): # This converts from iproc-then-ivar to ivar-then-iproc ordering # Subtract off part of previous procs # Then add all variables on previous procs # Then all previous variables on this proc # - np.sum(out_sizes[:iproc, idx_out]) # + np.sum(out_sizes[:iproc, :]) # + np.sum(out_sizes[iproc, :idx_out]) # + inds offset = offsets_out[iproc, idx_out] - start output_inds[ on_iproc] = src_indices[on_iproc] + offset start = end # 2. Compute the input indices input_inds = np.arange(offsets_in[myproc, idx_in], offsets_in[myproc, idx_in] + sizes_in[myproc, idx_in], dtype=INT_DTYPE) # Now the indices are ready - input_inds, output_inds sub_in = abs_in[mypathlen:].split('.', 1)[0] isub = sub_inds[sub_in] fwd_xfer_in[isub].append(input_inds) fwd_xfer_out[isub].append(output_inds) if rev: sub_out = abs_out[mypathlen:].split('.', 1)[0] isub = sub_inds[sub_out] rev_xfer_in[isub].append(input_inds) rev_xfer_out[isub].append(output_inds) transfers[vec_name] = {} for isub in range(nsub_allprocs): fwd_xfer_in[isub] = _merge(fwd_xfer_in[isub]) fwd_xfer_out[isub] = _merge(fwd_xfer_out[isub]) if rev: rev_xfer_in[isub] = _merge(rev_xfer_in[isub]) rev_xfer_out[isub] = _merge(rev_xfer_out[isub]) xfer_in = np.concatenate(fwd_xfer_in) xfer_out = np.concatenate(fwd_xfer_out) out_vec = vectors['output'][vec_name] xfer_all = PETScTransfer(vectors['input'][vec_name], out_vec, xfer_in, xfer_out, group.comm) transfers[vec_name]['fwd', None] = xfer_all if rev: transfers[vec_name]['rev', None] = xfer_all for isub in range(nsub_allprocs): transfers[vec_name]['fwd', isub] = PETScTransfer( vectors['input'][vec_name], vectors['output'][vec_name], fwd_xfer_in[isub], fwd_xfer_out[isub], group.comm) if rev: transfers[vec_name]['rev', isub] = PETScTransfer( vectors['input'][vec_name], vectors['output'][vec_name], rev_xfer_in[isub], rev_xfer_out[isub], group.comm) if group._use_derivatives: transfers['nonlinear'] = transfers['linear']
def _setup_transfers(group): """ Compute all transfers that are owned by our parent group. Parameters ---------- group : <Group> Parent group. """ rev = group._mode != 'fwd' for subsys in group._subgroups_myproc: subsys._setup_transfers() abs2meta_in = group._var_abs2meta['input'] abs2meta_out = group._var_abs2meta['output'] allprocs_abs2meta_out = group._var_allprocs_abs2meta['output'] myproc = group.comm.rank transfers = group._transfers = {} vectors = group._vectors offsets = group._get_var_offsets() mypathlen = len(group.pathname) + 1 if group.pathname else 0 allsubs = group._subsystems_allprocs # Initialize empty lists for the transfer indices xfer_in = [] xfer_out = [] fwd_xfer_in = defaultdict(list) fwd_xfer_out = defaultdict(list) if rev: rev_xfer_in = defaultdict(list) rev_xfer_out = defaultdict(list) allprocs_abs2idx = group._var_allprocs_abs2idx sizes_in = group._var_sizes['input'] sizes_out = group._var_sizes['output'] offsets_in = offsets['input'] offsets_out = offsets['output'] # Loop through all connections owned by this system for abs_in, abs_out in group._conn_abs_in2out.items(): # Only continue if the input exists on this processor if abs_in in abs2meta_in: # Get meta meta_in = abs2meta_in[abs_in] meta_out = allprocs_abs2meta_out[abs_out] idx_in = allprocs_abs2idx[abs_in] idx_out = allprocs_abs2idx[abs_out] # Read in and process src_indices src_indices = meta_in['src_indices'] if src_indices is None: owner = group._owning_rank[abs_out] # if the input is larger than the output on a single proc, we have # to just loop over the procs in the same way we do when src_indices # is defined. if meta_in['size'] > sizes_out[owner, idx_out]: src_indices = np.arange(meta_in['size'], dtype=INT_DTYPE) else: src_indices = src_indices.shaped_array() # 1. Compute the output indices # NOTE: src_indices are relative to a single, possibly distributed variable, # while the output_inds that we compute are relative to the full distributed # array that contains all local variables from each rank stacked in rank order. if src_indices is None: if meta_out['distributed']: # input in this case is non-distributed (else src_indices would be # defined by now). dist output to serial input conns w/o src_indices # are not allowed. raise RuntimeError( f"{group.msginfo}: Can't connect distributed output " f"'{abs_out}' to serial input '{abs_in}' without " "declaring src_indices.") else: rank = myproc if abs_out in abs2meta_out else owner offset = offsets_out[rank, idx_out] output_inds = np.arange(offset, offset + meta_in['size'], dtype=INT_DTYPE) else: output_inds = np.zeros(src_indices.size, INT_DTYPE) start = end = 0 for iproc in range(group.comm.size): end += sizes_out[iproc, idx_out] if start == end: continue # The part of src on iproc on_iproc = np.logical_and(start <= src_indices, src_indices < end) if np.any(on_iproc): # This converts from iproc-then-ivar to ivar-then-iproc ordering # Subtract off part of previous procs # Then add all variables on previous procs # Then all previous variables on this proc # - np.sum(out_sizes[:iproc, idx_out]) # + np.sum(out_sizes[:iproc, :]) # + np.sum(out_sizes[iproc, :idx_out]) # + inds offset = offsets_out[iproc, idx_out] - start output_inds[ on_iproc] = src_indices[on_iproc] + offset start = end # 2. Compute the input indices input_inds = np.arange(offsets_in[myproc, idx_in], offsets_in[myproc, idx_in] + sizes_in[myproc, idx_in], dtype=INT_DTYPE) # Now the indices are ready - input_inds, output_inds sub_in = abs_in[mypathlen:].split('.', 1)[0] fwd_xfer_in[sub_in].append(input_inds) fwd_xfer_out[sub_in].append(output_inds) if rev: sub_out = abs_out[mypathlen:].split('.', 1)[0] rev_xfer_in[sub_out].append(input_inds) rev_xfer_out[sub_out].append(output_inds) else: # not a local input but still need entries in the transfer dicts to # avoid hangs sub_in = abs_in[mypathlen:].split('.', 1)[0] fwd_xfer_in[ sub_in] # defaultdict will create an empty list there fwd_xfer_out[sub_in] if rev: sub_out = abs_out[mypathlen:].split('.', 1)[0] rev_xfer_in[sub_out] rev_xfer_out[sub_out] for sname, inds in fwd_xfer_in.items(): fwd_xfer_in[sname] = _merge(inds) fwd_xfer_out[sname] = _merge(fwd_xfer_out[sname]) if rev: for sname, inds in rev_xfer_out.items(): rev_xfer_in[sname] = _merge(rev_xfer_in[sname]) rev_xfer_out[sname] = _merge(inds) if fwd_xfer_in: xfer_in = np.concatenate(list(fwd_xfer_in.values())) xfer_out = np.concatenate(list(fwd_xfer_out.values())) else: xfer_in = xfer_out = np.zeros(0, dtype=INT_DTYPE) out_vec = vectors['output']['nonlinear'] xfer_all = PETScTransfer(vectors['input']['nonlinear'], out_vec, xfer_in, xfer_out, group.comm) transfers['fwd'] = xfwd = {} xfwd[None] = xfer_all if rev: transfers['rev'] = xrev = {} xrev[None] = xfer_all for sname, inds in fwd_xfer_in.items(): transfers['fwd'][sname] = PETScTransfer( vectors['input']['nonlinear'], vectors['output']['nonlinear'], inds, fwd_xfer_out[sname], group.comm) if rev: for sname, inds in rev_xfer_out.items(): transfers['rev'][sname] = PETScTransfer( vectors['input']['nonlinear'], vectors['output']['nonlinear'], rev_xfer_in[sname], inds, group.comm)