:mod:`heat.tiling` ======================= .. py:module:: heat.core.tiling .. autoapi-nested-parse:: Tiling functions/classes. With these classes, you can classes you can address blocks of data in a DNDarray Module Contents --------------- .. py:class:: SplitTiles(arr: heat.core.dndarray.DNDarray) Initialize tiles with the tile divisions equal to the theoretical split dimensions in every dimension :param arr: Base array for which to create the tiles :type arr: DNDarray :ivar __DNDarray: the ``DNDarray`` associated with the tiles :vartype __DNDarray: DNDarray :ivar __lshape_map: map of the shapes of the local torch tensors of arr :vartype __lshape_map: torch.Tensor :ivar __tile_locations: locations of the tiles of ``arr`` :vartype __tile_locations: torch.Tensor :ivar __tile_ends_g: the global indices of the ends of the tiles :vartype __tile_ends_g: torch.Tensor :ivar __tile_dims: the dimensions of all of the tiles :vartype __tile_dims: torch.Tensor .. rubric:: Examples >>> a = ht.zeros( ... ( ... 10, ... 11, ... ), ... split=None, ... ) >>> a.create_split_tiles() >>> print(a.tiles.tile_ends_g) [0/2] tensor([[ 4, 7, 10], [0/2] [ 4, 8, 11]], dtype=torch.int32) [1/2] tensor([[ 4, 7, 10], [1/2] [ 4, 8, 11]], dtype=torch.int32) [2/2] tensor([[ 4, 7, 10], [2/2] [ 4, 8, 11]], dtype=torch.int32) >>> print(a.tiles.tile_locations) [0/2] tensor([[0, 0, 0], [0/2] [0, 0, 0], [0/2] [0, 0, 0]], dtype=torch.int32) [1/2] tensor([[1, 1, 1], [1/2] [1, 1, 1], [1/2] [1, 1, 1]], dtype=torch.int32) [2/2] tensor([[2, 2, 2], [2/2] [2, 2, 2], [2/2] [2, 2, 2]], dtype=torch.int32) >>> a = ht.zeros((10, 11), split=1) >>> a.create_split_tiles() >>> print(a.tiles.tile_ends_g) [0/2] tensor([[ 4, 7, 10], [0/2] [ 4, 8, 11]], dtype=torch.int32) [1/2] tensor([[ 4, 7, 10], [1/2] [ 4, 8, 11]], dtype=torch.int32) [2/2] tensor([[ 4, 7, 10], [2/2] [ 4, 8, 11]], dtype=torch.int32) >>> print(a.tiles.tile_locations) [0/2] tensor([[0, 1, 2], [0/2] [0, 1, 2], [0/2] [0, 1, 2]], dtype=torch.int32) [1/2] tensor([[0, 1, 2], [1/2] [0, 1, 2], [1/2] [0, 1, 2]], dtype=torch.int32) [2/2] tensor([[0, 1, 2], [2/2] [0, 1, 2], [2/2] [0, 1, 2]], dtype=torch.int32) .. attribute:: __DNDarray .. attribute:: __lshape_map .. attribute:: __tile_locations .. attribute:: __tile_ends_g .. attribute:: __tile_dims .. role:: raw-html(raw) :format: html .. method:: set_tile_locations(split: int, tile_dims: torch.Tensor, arr: heat.core.dndarray.DNDarray) -> torch.Tensor Create a `torch.Tensor` which contains the locations of the tiles of ``arr`` for the given split :param split: Target split dimension. Does not need to be equal to ``arr.split`` :type split: int :param tile_dims: Tensor containing the sizes of the each tile :type tile_dims: torch.Tensor :param arr: Array for which the tiles are being created for :type arr: DNDarray .. method:: __getitem__(key: Union[int, slice, Tuple[Union[int, slice], Ellipsis]]) -> torch.Tensor Getitem function for getting tiles. Returns the tile which is specified is returned, but only on the process which it resides :param key: Key which identifies the tile/s to get :type key: int or Tuple or Slice .. rubric:: Examples >>> test = torch.arange(np.prod([i + 6 for i in range(2)])).reshape( ... [i + 6 for i in range(2)] ... ) >>> a = ht.array(test, split=0).larray [0/2] tensor([[ 0., 1., 2., 3., 4., 5., 6.], [0/2] [ 7., 8., 9., 10., 11., 12., 13.]]) [1/2] tensor([[14., 15., 16., 17., 18., 19., 20.], [1/2] [21., 22., 23., 24., 25., 26., 27.]]) [2/2] tensor([[28., 29., 30., 31., 32., 33., 34.], [2/2] [35., 36., 37., 38., 39., 40., 41.]]) >>> a.create_split_tiles() >>> a.tiles[:2, 2] [0/2] tensor([[ 5., 6.], [0/2] [12., 13.]]) [1/2] tensor([[19., 20.], [1/2] [26., 27.]]) [2/2] None >>> a = ht.array(test, split=1) >>> a.create_split_tiles() >>> a.tiles[1] [0/2] tensor([[14., 15., 16.], [0/2] [21., 22., 23.]]) [1/2] tensor([[17., 18.], [1/2] [24., 25.]]) [2/2] tensor([[19., 20.], [2/2] [26., 27.]]) .. method:: __get_tile_slices(key: Union[int, slice, Tuple[Union[int, slice], Ellipsis]]) -> Tuple[slice, Ellipsis] Create and return slices to convert a key from the tile indices to the normal indices .. method:: get_tile_size(key: Union[int, slice, Tuple[Union[int, slice], Ellipsis]]) -> Tuple[int, Ellipsis] Get the size of a tile or tiles indicated by the given key :param key: which tiles to get :type key: int or slice or tuple .. method:: __setitem__(key: Union[int, slice, Tuple[Union[int, slice], Ellipsis]], value: Union[int, float, torch.Tensor]) -> None Set the values of a tile :param key: Key which identifies the tile/s to get :type key: int or Tuple or Slice :param value: Value to be set on the tile :type value: int or torch.Tensor .. rubric:: Examples see getitem function for this class .. method:: get_subarray_params(from_axis: int, to_axis: int) -> List[Tuple[List[int], List[int], List[int]]] Create subarray types of the local array along a new split axis. For use with Alltoallw. Based on the work by Dalcin et al. (https://arxiv.org/abs/1804.09536) Return type is a list of tuples, each tuple containing the shape of the local array, the shape of the subarray, and the start index of the subarray. :param from_axis: Current split axis of global array. :type from_axis: int :param to_axis: New split axis of of subarrays array. :type to_axis: int .. py:class:: SquareDiagTiles(arr: heat.core.dndarray.DNDarray, tiles_per_proc: int = 2) Generate the tile map and the other objects which may be useful. The tiles generated here are based of square tiles along the diagonal. The size of these tiles along the diagonal dictate the divisions across all processes. If ``gshape[0]>>gshape[1]`` then there will be extra tiles generated below the diagonal. If ``gshape[0]`` is close to ``gshape[1]``, then the last tile (as well as the other tiles which correspond with said tile) will be extended to cover the whole array. However, extra tiles are not generated above the diagonal in the case that ``gshape[0]< [rank, row size, column size]`` Tensor filled with the shapes of the local tensors :vartype __lshape_map: torch.Tensor :ivar __tile_map: ``units -> row, column, start index in each direction, process`` Tensor filled with the global indices of the generated tiles :vartype __tile_map: torch.Tensor :ivar __row_per_proc_list: List is length of the number of processes, each element has the number of tile rows on the process whos rank equals the index :vartype __row_per_proc_list: List .. warning:: The generation of these tiles may unbalance the original ``DNDarray``! .. rubric:: Notes This tiling scheme is intended for use with the :func:`~heat.core.linalg.qr.qr` function. .. attribute:: __DNDarray .. attribute:: __col_per_proc_list .. attribute:: __lshape_map .. attribute:: __last_diag_pr .. attribute:: __row_per_proc_list .. attribute:: __tile_map .. attribute:: __row_inds .. attribute:: __col_inds .. role:: raw-html(raw) :format: html .. method:: __adjust_cols_sp1_m_ls_n(arr: heat.core.dndarray.DNDarray, col_per_proc_list: List[int, Ellipsis], last_diag_pr: int, col_inds: List[int, Ellipsis], lshape_map: torch.Tensor) -> None Add more columns after the diagonal ends if ``m None Need to adjust the size of last row if ``arr.split==0`` and the diagonal ends before the last tile. This should only be run if ``arr,split==0`` and ``last_diag_pr None If the split is 0 and the number of tiles per proc is 1 then the local data may need to be redistributed to fit the full diagonal on as many processes as possible. If there is a process where there is only 1 element, this function will adjust the ``lshape_map`` then redistribute ``arr`` so that there is not a single diagonal element on one process .. method:: __create_cols(arr: heat.core.dndarray.DNDarray, lshape_map: torch.Tensor, tiles_per_proc: int) -> Tuple[torch.Tensor, List[int, Ellipsis], List[int, Ellipsis], torch.Tensor] Calculates the last diagonal process, then creates a list of the number of tile columns per process, then calculates the starting indices of the columns. Also returns the number of tile columns. :param arr: DNDarray for which to find the tile columns for :type arr: DNDarray :param lshape_map: The map of the local shapes (for more info see: :func:`~heat.core.dndarray.DNDarray.create_lshape_map`) :type lshape_map: torch.Tensor :param tiles_per_proc: The number of divisions per process :type tiles_per_proc: int .. method:: __def_end_row_inds_sp0_m_ge_n(arr: heat.core.dndarray.DNDarray, row_inds: List[int, Ellipsis], last_diag_pr: int, tiles_per_proc: int, lshape_map: torch.Tensor) -> None Adjust the rows on the processes which are greater than the last diagonal processs to have rows which are chunked evenly into ``tiles_per_proc`` rows. .. method:: __last_tile_row_adjust_sp1(arr: heat.core.dndarray.DNDarray, row_inds: List[int, Ellipsis]) -> None Add extra row/s if there is space below the diagonal (``split=1``) .. method:: get_start_stop(key: Union[int, slice, Tuple[int, slice, Ellipsis]]) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] Returns the start and stop indices in form of ``(dim0 start, dim0 stop, dim1 start, dim1 stop)`` which correspond to the tile/s which corresponds to the given key. The key MUST use global indices. :param key: Indices to select the tile STRIDES ARE NOT ALLOWED, MUST BE GLOBAL INDICES :type key: int or Tuple or List or slice .. rubric:: Examples >>> a = ht.zeros((12, 10), split=0) >>> a_tiles = ht.tiling.SquareDiagTiles(a, tiles_per_proc=2) # type: tiling.SquareDiagTiles >>> print(a_tiles.get_start_stop(key=(slice(0, 2), 2))) [0/1] (tensor(0), tensor(6), tensor(6), tensor(8)) [1/1] (tensor(0), tensor(6), tensor(6), tensor(8)) >>> print(a_tiles.get_start_stop(key=(0, 2))) [0/1] (tensor(0), tensor(3), tensor(6), tensor(8)) [1/1] (tensor(0), tensor(3), tensor(6), tensor(8)) >>> print(a_tiles.get_start_stop(key=2)) [0/1] (tensor(0), tensor(2), tensor(0), tensor(10)) [1/1] (tensor(0), tensor(2), tensor(0), tensor(10)) >>> print(a_tiles.get_start_stop(key=(3, 3))) [0/1] (tensor(2), tensor(6), tensor(8), tensor(10)) [1/1] (tensor(2), tensor(6), tensor(8), tensor(10)) .. method:: __getitem__(key: Union[int, slice, Tuple[int, slice, Ellipsis]]) -> torch.Tensor Returns a local selection of the DNDarray corresponding to the tile/s desired Standard getitem function for the tiles. The returned item is a view of the original DNDarray, operations which are done to this view will change the original array. **STRIDES ARE NOT AVAILABLE, NOR ARE CROSS-SPLIT SLICES** :param key: indices of the tile/s desired :type key: int, slice, tuple .. rubric:: Examples >>> a = ht.zeros((12, 10), split=0) >>> a_tiles = tiling.SquareDiagTiles(a, tiles_per_proc=2) # type: tiling.SquareDiagTiles >>> print(a_tiles[2, 3]) [0/1] None [1/1] tensor([[0., 0.], [1/1] [0., 0.]]) >>> print(a_tiles[2]) [0/1] None [1/1] tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], [1/1] [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]]) >>> print(a_tiles[0:2, 1]) [0/1] tensor([[0., 0., 0.], [0/1] [0., 0., 0.], [0/1] [0., 0., 0.], [0/1] [0., 0., 0.], [0/1] [0., 0., 0.], [0/1] [0., 0., 0.]]) [1/1] None .. method:: local_get(key: Union[int, slice, Tuple[int, slice, Ellipsis]]) -> torch.Tensor Returns the local tile/s corresponding to the key given Getitem routing using local indices, converts to global indices then uses getitem :param key: Indices of the tile/s desired. If the stop index of a slice is larger than the end will be adjusted to the maximum allowed :type key: int, slice, tuple, list .. rubric:: Examples See local_set function. .. method:: local_set(key: Union[int, slice, Tuple[int, slice, Ellipsis]], value: Union[int, float, torch.Tensor]) Setitem routing to set data to a local tile (using local indices) :param key: Indices of the tile/s desired If the stop index of a slice is larger than the end will be adjusted to the maximum allowed :type key: int or slice or Tuple[int,...] :param value: Data to be written to the tile :type value: torch.Tensor or int or float .. rubric:: Examples >>> a = ht.zeros((11, 10), split=0) >>> a_tiles = tiling.SquareDiagTiles(a, tiles_per_proc=2) # type: tiling.SquareDiagTiles >>> local = a_tiles.local_get(key=slice(None)) >>> a_tiles.local_set( ... key=slice(None), value=torch.arange(local.numel()).reshape(local.shape) ... ) >>> print(a.larray) [0/1] tensor([[ 0., 1., 2., 3., 4., 5., 6., 7., 8., 9.], [0/1] [10., 11., 12., 13., 14., 15., 16., 17., 18., 19.], [0/1] [20., 21., 22., 23., 24., 25., 26., 27., 28., 29.], [0/1] [30., 31., 32., 33., 34., 35., 36., 37., 38., 39.], [0/1] [40., 41., 42., 43., 44., 45., 46., 47., 48., 49.], [0/1] [50., 51., 52., 53., 54., 55., 56., 57., 58., 59.]]) [1/1] tensor([[ 0., 1., 2., 3., 4., 5., 6., 7., 8., 9.], [1/1] [10., 11., 12., 13., 14., 15., 16., 17., 18., 19.], [1/1] [20., 21., 22., 23., 24., 25., 26., 27., 28., 29.], [1/1] [30., 31., 32., 33., 34., 35., 36., 37., 38., 39.], [1/1] [40., 41., 42., 43., 44., 45., 46., 47., 48., 49.]]) >>> a.lloc[:] = 0 >>> a_tiles.local_set(key=(0, 2), value=10) [0/1] tensor([[ 0., 0., 0., 0., 0., 0., 10., 10., 0., 0.], [0/1] [ 0., 0., 0., 0., 0., 0., 10., 10., 0., 0.], [0/1] [ 0., 0., 0., 0., 0., 0., 10., 10., 0., 0.], [0/1] [ 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], [0/1] [ 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], [0/1] [ 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]]) [1/1] tensor([[ 0., 0., 0., 0., 0., 0., 10., 10., 0., 0.], [1/1] [ 0., 0., 0., 0., 0., 0., 10., 10., 0., 0.], [1/1] [ 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], [1/1] [ 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], [1/1] [ 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]]) >>> a_tiles.local_set(key=(slice(None), 1), value=10) [0/1] tensor([[ 0., 0., 0., 10., 10., 10., 0., 0., 0., 0.], [0/1] [ 0., 0., 0., 10., 10., 10., 0., 0., 0., 0.], [0/1] [ 0., 0., 0., 10., 10., 10., 0., 0., 0., 0.], [0/1] [ 0., 0., 0., 10., 10., 10., 0., 0., 0., 0.], [0/1] [ 0., 0., 0., 10., 10., 10., 0., 0., 0., 0.], [0/1] [ 0., 0., 0., 10., 10., 10., 0., 0., 0., 0.]]) [1/1] tensor([[ 0., 0., 0., 10., 10., 10., 0., 0., 0., 0.], [1/1] [ 0., 0., 0., 10., 10., 10., 0., 0., 0., 0.], [1/1] [ 0., 0., 0., 10., 10., 10., 0., 0., 0., 0.], [1/1] [ 0., 0., 0., 10., 10., 10., 0., 0., 0., 0.], [1/1] [ 0., 0., 0., 10., 10., 10., 0., 0., 0., 0.]]) .. method:: local_to_global(key: Union[int, slice, Tuple[int, slice, Ellipsis]], rank: int) -> Tuple[int, slice, Ellipsis] Convert local indices to global indices :param key: Indices of the tile/s desired. If the stop index of a slice is larger than the end will be adjusted to the maximum allowed :type key: int or slice or Tuple or List :param rank: Process rank :type rank: int .. rubric:: Examples >>> a = ht.zeros((11, 10), split=0) >>> a_tiles = tiling.SquareDiagTiles(a, tiles_per_proc=2) # type: tiling.SquareDiagTiles >>> rank = a.comm.rank >>> print(a_tiles.local_to_global(key=(slice(None), 1), rank=rank)) [0/1] (slice(0, 2, None), 1) [1/1] (slice(2, 4, None), 1) >>> print(a_tiles.local_to_global(key=(0, 2), rank=0)) [0/1] (0, 2) [1/1] (0, 2) >>> print(a_tiles.local_to_global(key=(0, 2), rank=1)) [0/1] (2, 2) [1/1] (2, 2) .. method:: match_tiles(tiles_to_match: SquareDiagTiles) -> None Function to match the tile sizes of another tile map :param tiles_to_match: The tiles which should be matched by the current tiling scheme :type tiles_to_match: SquareDiagTiles .. rubric:: Notes This function overwrites most, if not all, of the elements of this class. Intended for use with the Q matrix, to match the tiling of a/R. For this to work properly it is required that the 0th dim of both matrices is equal .. method:: __setitem__(key: Union[int, slice, Tuple[int, slice, Ellipsis]], value: Union[int, float, torch.Tensor]) -> None Item setter, uses the torch item setter and the getitem routines to set the values of the original array (arr in __init__) :param key: Tile indices to identify the target tiles :type key: int or slice or Tuple[int,...] :param value: Values to be set :type value: int or torch.Tensor .. rubric:: Example >>> a = ht.zeros((12, 10), split=0) >>> a_tiles = tiling.SquareDiagTiles(a, tiles_per_proc=2) # type: tiling.SquareDiagTiles >>> a_tiles[0:2, 2] = 11 >>> a_tiles[0, 0] = 22 >>> a_tiles[2] = 33 >>> a_tiles[3, 3] = 44 >>> print(a.larray) [0/1] tensor([[22., 22., 22., 0., 0., 0., 11., 11., 0., 0.], [0/1] [22., 22., 22., 0., 0., 0., 11., 11., 0., 0.], [0/1] [22., 22., 22., 0., 0., 0., 11., 11., 0., 0.], [0/1] [ 0., 0., 0., 0., 0., 0., 11., 11., 0., 0.], [0/1] [ 0., 0., 0., 0., 0., 0., 11., 11., 0., 0.], [0/1] [ 0., 0., 0., 0., 0., 0., 11., 11., 0., 0.]]) [1/1] tensor([[33., 33., 33., 33., 33., 33., 33., 33., 33., 33.], [1/1] [33., 33., 33., 33., 33., 33., 33., 33., 33., 33.], [1/1] [ 0., 0., 0., 0., 0., 0., 0., 0., 44., 44.], [1/1] [ 0., 0., 0., 0., 0., 0., 0., 0., 44., 44.], [1/1] [ 0., 0., 0., 0., 0., 0., 0., 0., 44., 44.], [1/1] [ 0., 0., 0., 0., 0., 0., 0., 0., 44., 44.]])