heat.tiling

Tiling functions/classes. With these classes, you can classes you can address blocks of data in a DNDarray

Module Contents

class SplitTiles(arr: heat.core.dndarray.DNDarray)

Initialize tiles with the tile divisions equal to the theoretical split dimensions in every dimension

Parameters:

arr (DNDarray) – Base array for which to create the tiles

Variables:

• __DNDarray (DNDarray) – the DNDarray associated with the tiles

• __lshape_map (torch.Tensor) – map of the shapes of the local torch tensors of arr

• __tile_locations (torch.Tensor) – locations of the tiles of arr

• __tile_ends_g (torch.Tensor) – the global indices of the ends of the tiles

• __tile_dims (torch.Tensor) – the dimensions of all of the tiles

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)

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

Parameters:

• split (int) – Target split dimension. Does not need to be equal to arr.split

• tile_dims (torch.Tensor) – Tensor containing the sizes of the each tile

• arr (DNDarray) – Array for which the tiles are being created for

__getitem__(key: int | slice | Tuple[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

Parameters:

key (int or Tuple or Slice) – Key which identifies the tile/s to get

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.]])

__get_tile_slices(key: int | slice | Tuple[int | slice, Ellipsis]) → Tuple[slice, Ellipsis]

Create and return slices to convert a key from the tile indices to the normal indices

get_tile_size(key: int | slice | Tuple[int | slice, Ellipsis]) → Tuple[int, Ellipsis]

Get the size of a tile or tiles indicated by the given key

Parameters:

key (int or slice or tuple) – which tiles to get

__setitem__(key: int | slice | Tuple[int | slice, Ellipsis], value: int | float | torch.Tensor) → None

Set the values of a tile

Parameters:

• key (int or Tuple or Slice) – Key which identifies the tile/s to get

• value (int or torch.Tensor) – Value to be set on the tile

Examples

see getitem function for this class

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]<<gshape[1].

Parameters:

• arr (DNDarray) – The array to be tiled

• tiles_per_proc (int, optional) – The number of divisions per process Default: 2

Variables:

• __col_per_proc_list (List) – List is length of the number of processes, each element has the number of tile columns on the process whos rank equals the index

• __DNDarray (DNDarray) – The whole DNDarray

• __lshape_map (torch.Tensor) – unit -> [rank, row size, column size] Tensor filled with the shapes of the local tensors

• __tile_map (torch.Tensor) – units -> row, column, start index in each direction, process Tensor filled with the global indices of the generated tiles

• __row_per_proc_list (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

Warning

The generation of these tiles may unbalance the original DNDarray!

Notes

This tiling scheme is intended for use with the qr() function.

__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<n and arr.split==1

__adjust_last_row_sp0_m_ge_n(arr: heat.core.dndarray.DNDarray, lshape_map: torch.Tensor, last_diag_pr: int, row_inds: List[int, Ellipsis], row_per_proc_list: List[int, Ellipsis], tile_columns: int) → 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<arr.comm.size-1.

__adjust_lshape_sp0_1tile(arr: heat.core.dndarray.DNDarray, col_inds: List[int, Ellipsis], lshape_map: torch.Tensor, tiles_per_proc: int) → 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

__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.

Parameters:

• arr (DNDarray) – DNDarray for which to find the tile columns for

• lshape_map (torch.Tensor) – The map of the local shapes (for more info see: create_lshape_map())

• tiles_per_proc (int) – The number of divisions per process

__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.

__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)

get_start_stop(key: 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.

Parameters:

key (int or Tuple or List or slice) – Indices to select the tile STRIDES ARE NOT ALLOWED, MUST BE GLOBAL INDICES

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))

__getitem__(key: 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

Parameters:

key (int, slice, tuple) – indices of the tile/s desired

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

local_get(key: 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

Parameters:

key (int, slice, tuple, list) – 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

Examples

See local_set function.

local_set(key: int | slice | Tuple[int, slice, Ellipsis], value: int | float | torch.Tensor)

Setitem routing to set data to a local tile (using local indices)

Parameters:

• key (int or slice or Tuple[int,...]) – 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

• value (torch.Tensor or int or float) – Data to be written to the tile

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.]])

local_to_global(key: int | slice | Tuple[int, slice, Ellipsis], rank: int) → Tuple[int, slice, Ellipsis]

Convert local indices to global indices

Parameters:

• key (int or slice or Tuple or List) – 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

• rank (int) – Process rank

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)

match_tiles(tiles_to_match: SquareDiagTiles) → None

Function to match the tile sizes of another tile map

Parameters:

tiles_to_match (SquareDiagTiles) – The tiles which should be matched by the current tiling scheme

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

__setitem__(key: int | slice | Tuple[int, slice, Ellipsis], value: 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__)

Parameters:

• key (int or slice or Tuple[int,...]) – Tile indices to identify the target tiles

• value (int or torch.Tensor) – Values to be set

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.]])