warp.tile_store_indexed#

warp.tile_store_indexed(
a: Array[Any],
indices: Tile[int32, tuple[int]],
t: Tile[Any, tuple[int, ...]],
offset: tuple[int, ...] = ...,
axis: int32 = 0,
) → None#
  • Kernel
  • Differentiable

Store a tile to a global memory array, scattering along one axis through a 1D tile of indices.

Cooperative operation: every thread of the block must reach the call. Element c of t is written to 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 destination index is negative or past the end of a is skipped, so -1 can be used to discard a slice.

The selected elements of a are overwritten. Each destination must be selected by at most one element — duplicate indices race. axis must be a compile-time constant. The backward pass accumulates gradients at the written destinations into the adjoint of t, then clears those entries from a.grad.

Parameters:

  • a – The destination array in global memory

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

  • t – The source tile to store data from, must have the same data type and number of dimensions as the destination array, and along axis the same number of elements as the indices tile

  • offset – Offset in the destination array, 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

Example

This example writes the rows of a tile to the even-numbered rows of a 2D array.

TILE_M, TILE_N = 2, 4
TILE_THREADS = 4

@wp.kernel
def scatter_even_rows(x: wp.array2d[float], y: wp.array2d[float]):
    t = wp.tile_load(x, shape=(TILE_M, TILE_N))
    # tile row k is written to row 2*k of `y`
    rows = wp.tile_arange(TILE_M, dtype=int) * 2
    wp.tile_store_indexed(y, indices=rows, t=t, axis=0)

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

[[1. 2. 3. 4.]
 [0. 0. 0. 0.]
 [5. 6. 7. 8.]
 [0. 0. 0. 0.]]