warp.tile_atomic_add#

warp.tile_atomic_add(
a: Array[Any],
t: Tile[Any, tuple[int, ...]],
offset: tuple[int, ...] = ...,
bounds_check: bool = True,
) → Tile[Any, tuple[int, ...]]#
  • Kernel
  • Differentiable

Atomically add a tile onto the array a and return the values it replaced.

This is a cooperative operation: the threads of the block divide the work between them, so every thread must reach the call. Element (i, j, ...) of t is added atomically to a[offset[0] + i, offset[1] + j, ...] and the destination’s previous value is placed in the returned tile. No barrier is issued by the call itself.

Only the individual element updates are atomic. Concurrent updates from other threads or blocks are interleaved in an unspecified order, so the returned values — and, for floating-point types, the rounding of the accumulated result — are not reproducible. For Warp struct elements, only fields whose underlying scalar type supports atomic addition are updated. Boolean, narrow-integer, array, and other non-atomic fields remain unchanged, although their previous values are still present in the returned tile.

In a backward pass the gradients of the updated region of a are accumulated into the adjoint of t and left in place in a.grad; the adjoint of the returned tile is not propagated.

Parameters:

  • a – Array in global memory, must have the same dtype as the input tile. Its underlying scalar type must be one that supports atomic addition: int32, uint32, int64, uint64, float16, bfloat16, float32, or float64.

  • t – Source tile to add to the destination array

  • offset – Offset in the destination array, one value per dimension of a; may be a runtime value.

  • bounds_check – Whether to treat a destination coordinate at or past the array’s upper extent on any axis as out of bounds; such updates are skipped. Must be a compile-time constant.

Returns:

A tile with the same shape, data type and storage as t, holding the value each destination element had before the addition. Passing a shared t therefore allocates a second shared-memory tile for the result.

Example

TILE_THREADS = 2

@wp.kernel
def accumulate(x: wp.array2d[float], totals: wp.array2d[float], previous: wp.array2d[float]):
    t = wp.tile_load(x, shape=(2, 2))
    # `p` holds what `totals` contained before the addition
    p = wp.tile_atomic_add(totals, t)
    wp.tile_store(previous, p)

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

[[1. 3.]
 [5. 7.]]
[[0. 1.]
 [2. 3.]]
warp.tile_atomic_add(
a: Array[Any],
t: Tile[Any, tuple[int, ...]],
offset: int32 = ...,
bounds_check: bool = True,
) → Tile[Any, tuple[int, ...]]
  • Kernel
  • Differentiable

Atomically add a 1D tile onto the 1D array a and return the values it replaced.

Overload for a scalar offset, equivalent to passing a one-element tuple. For the full contract and a usage example, see the overload that takes a tuple-valued offset argument.