def _cart_create(pv: 'ProgramVisitor', sdfg: SDFG, state: SDFGState, dims: ShapeType): """ Creates a process-grid and adds it to the DaCe program. The process-grid is implemented with [MPI_Cart_create](https://www.mpich.org/static/docs/latest/www3/MPI_Cart_create.html). :param dims: Shape of the process-grid (see `dims` parameter of `MPI_Cart_create`), e.g., [2, 3, 3]. :return: Name of the new process-grid descriptor. """ pgrid_name = sdfg.add_pgrid(dims) # Dummy tasklet adds MPI variables to the program's state. from dace.libraries.mpi import Dummy tasklet = Dummy(pgrid_name, [ f'MPI_Comm {pgrid_name}_comm;', f'MPI_Group {pgrid_name}_group;', f'int {pgrid_name}_coords[{len(dims)}];', f'int {pgrid_name}_dims[{len(dims)}];', f'int {pgrid_name}_rank;', f'int {pgrid_name}_size;', f'bool {pgrid_name}_valid;', ]) state.add_node(tasklet) # Pseudo-writing to a dummy variable to avoid removal of Dummy node by transformations. _, scal = sdfg.add_scalar(pgrid_name, dace.int32, transient=True) wnode = state.add_write(pgrid_name) state.add_edge(tasklet, '__out', wnode, None, Memlet.from_array(pgrid_name, scal)) return pgrid_name
def _cart_sub(pv: 'ProgramVisitor', sdfg: SDFG, state: SDFGState, parent_grid: str, color: Sequence[Union[Integral, bool]], exact_grid: RankType = None): """ Partitions the `parent_grid` to lower-dimensional sub-grids and adds them to the DaCe program. The sub-grids are implemented with [MPI_Cart_sub](https://www.mpich.org/static/docs/latest/www3/MPI_Cart_sub.html). :param parent_grid: Parent process-grid (similar to the `comm` parameter of `MPI_Cart_sub`). :param color: The i-th entry specifies whether the i-th dimension is kept in the sub-grid or is dropped (see `remain_dims` input of `MPI_Cart_sub`). :param exact_grid: [DEVELOPER] If set then, out of all the sub-grids created, only the one that contains the rank with id `exact_grid` will be utilized for collective communication. :return: Name of the new sub-grid descriptor. """ pgrid_name = sdfg.add_pgrid(parent_grid=parent_grid, color=color, exact_grid=exact_grid) # Count sub-grid dimensions. pgrid_ndims = sum([bool(c) for c in color]) # Dummy tasklet adds MPI variables to the program's state. from dace.libraries.mpi import Dummy tasklet = Dummy(pgrid_name, [ f'MPI_Comm {pgrid_name}_comm;', f'MPI_Group {pgrid_name}_group;', f'int {pgrid_name}_coords[{pgrid_ndims}];', f'int {pgrid_name}_dims[{pgrid_ndims}];', f'int {pgrid_name}_rank;', f'int {pgrid_name}_size;', f'bool {pgrid_name}_valid;', ]) state.add_node(tasklet) # Pseudo-writing to a dummy variable to avoid removal of Dummy node by transformations. _, scal = sdfg.add_scalar(pgrid_name, dace.int32, transient=True) wnode = state.add_write(pgrid_name) state.add_edge(tasklet, '__out', wnode, None, Memlet.from_array(pgrid_name, scal)) return pgrid_name