warp.tile_load_indexed#

warp.tile_load_indexed(
a: Array[Any],
indices: Tile[int32, tuple[int]],
shape: tuple[int, ...],
offset: tuple[int, ...] = ...,
axis: int32 = 0,
storage: str = 'register',
) → Tile[Any, tuple[int, ...]]#
  • Kernel
  • Differentiable

Load a tile from a global memory array, gathering along one axis through a 1D tile of indices.

Cooperative operation: every thread of the block must reach the call. Tile element c is read from a at offset[d] + c[d] along every dimension d other than axis, and at offset[axis] + indices[c[axis]] along axis. Every coordinate is checked against both the lower and upper array bounds: an element whose source index is negative or past the end of a reads as zero, so -1 can be used as a padding sentinel without a physical zero row.

shape, axis, and storage must be compile-time constants. In a backward pass the adjoint of the returned tile is atomically accumulated into a.grad at the same gathered locations.

Parameters:

  • a – The source array in global memory

  • indices – A 1D tile of int32 indices into a along axis. It must hold exactly shape[axis] values and is always placed in shared memory (a register tile passed here is promoted).

  • shape – Shape of the tile to load, must have the same number of dimensions as a, and along axis the same number of elements as the indices tile

  • offset – Offset in the source array to begin reading from, one value per dimension of a; the entry for axis is added to each index; may be a runtime value.

  • axis – Axis of a that the indices refer to

  • storage – The storage location for the tile: "register" for registers or "shared" for shared memory.

Returns:

A tile with shape as specified and data type the same as the source array.

Example

This example gathers the even-numbered rows of a 2D array.

TILE_M, TILE_N = 2, 4
TILE_THREADS = 4

@wp.kernel
def gather_even_rows(x: wp.array2d[float], y: wp.array2d[float]):
    # gather rows 0, 2, 4, ... of `x`
    rows = wp.tile_arange(TILE_M, dtype=int) * 2
    t = wp.tile_load_indexed(x, indices=rows, shape=(TILE_M, TILE_N), axis=0)
    wp.tile_store(y, t)

x = wp.array(np.arange(1, 17, dtype=np.float32).reshape(4, 4), dtype=float)
y = wp.zeros((2, 4), dtype=float)
wp.launch_tiled(gather_even_rows, dim=1, inputs=[x], outputs=[y], block_dim=TILE_THREADS)
print(y.numpy())

[[ 1.  2.  3.  4.]
 [ 9. 10. 11. 12.]]