warp.tile_scatter_masked#

warp.tile_scatter_masked( a: Tile[Any, tuple[int, ...]], i: int32, value: Any, has_value: bool, ) → None#

Kernel: true
Python: false
Differentiable: true

Write a per-thread value into a shared-memory tile.

This is a cooperative operation, so every thread in the block must call it. Threads with has_value=True write value at index i; threads with nothing to write pass has_value=False. The writes are available to subsequent tile operations when the call returns.

Each index must be written by at most one thread per call; conflicting writes are undefined. Use tile_scatter_add() when several threads may target the same index.

Because the values come from individual threads, the result depends on block_dim and differs on CPU, which runs a single lane per block (see CPU Tile Semantics). In a backward pass the adjoint of value takes the tile’s gradient at i, which is then cleared.

Parameters:

a – Tile to write into; a register tile is promoted to shared memory.
i – Index of the element to write. Must be valid when has_value is True.
value – The value to write (must match the tile’s dtype).
has_value – Whether this thread should perform the write.

Example

@wp.kernel
def reverse_lanes(src: wp.array[int], dst: wp.array[int]):
    _block, i = wp.tid()
    t = wp.tile_zeros(shape=8, dtype=int, storage="shared")
    # a permutation: every slot is written by exactly one thread
    wp.tile_scatter_masked(t, 7 - i, src[i], True)
    wp.tile_store(dst, t)

src = wp.array(np.arange(1, 9), dtype=int)
dst = wp.zeros(8, dtype=int)
wp.launch_tiled(reverse_lanes, dim=1, inputs=[src], outputs=[dst], block_dim=8)
print(dst.numpy())

[8 7 6 5 4 3 2 1]

warp.tile_scatter_masked( a: Tile[Any, tuple[int, ...]], i: int32, j: int32, value: Any, has_value: bool, ) → None

Kernel: true
Python: false
Differentiable: true

Write a per-thread value into a 2D shared-memory tile.

Overload taking one index per tile dimension. For the full contract and a usage example, see the 1D overload that takes only i as its index.

warp.tile_scatter_masked( a: Tile[Any, tuple[int, ...]], i: int32, j: int32, k: int32, value: Any, has_value: bool, ) → None

Kernel: true
Python: false
Differentiable: true

Write a per-thread value into a 3D shared-memory tile.

Overload taking one index per tile dimension. For the full contract and a usage example, see the 1D overload that takes only i as its index.

warp.tile_scatter_masked( a: Tile[Any, tuple[int, ...]], i: int32, j: int32, k: int32, l: int32, value: Any, has_value: bool, ) → None

Kernel: true
Python: false
Differentiable: true

Write a per-thread value into a 4D shared-memory tile.

Overload taking one index per tile dimension. For the full contract and a usage example, see the 1D overload that takes only i as its index.