warp.tile_reduce#

warp.tile_reduce( op: Callable, a: Tile[Any, tuple[int, ...]], ) → Tile[Any, tuple[Literal[1]]]#

Kernel: true
Python: false
Differentiable: false

Apply a custom reduction operator across all elements of a tile.

This is a cooperative operation: the threads of the block divide the work between them, so every thread must reach the call.

Elements of a are combined in an unspecified order, so op must be associative and commutative for the result to be well defined; a non-commutative operator such as sub() gives different answers for different values of block_dim and on different devices.

Parameters:

op – Callable taking two elements of the tile’s dtype and returning one value of the same type. May be a Warp built-in such as add() or a user-defined function decorated with @wp.func.
a – The input tile. The operator, or one of its overloads, must accept the tile’s element type.

Returns:

A shape-(1,) register tile with the same dtype as a. Read the value with [0].

Example

@wp.kernel
def multiply_elements(x: wp.array[int], out: wp.array[int]):
    t = wp.tile_load(x, shape=4)
    wp.tile_store(out, wp.tile_reduce(wp.mul, t))

x = wp.array([3, 7, 4, 9], dtype=int)
out = wp.zeros(1, dtype=int)

wp.launch_tiled(multiply_elements, dim=1, inputs=[x], outputs=[out], block_dim=2)

print(out.numpy())

[756]

warp.tile_reduce( op: Callable, a: Tile[Scalar, tuple[int, ...]], axis: int32, ) → Tile[Scalar, tuple[int, ...]]

Kernel: true
Python: false
Differentiable: false

Apply a custom reduction operator along one axis of a tile.

This is a cooperative operation: the threads of the block divide the work between them, so every thread must reach the call.

Reduces over axis. Each slice is folded with op in an unspecified order, so op must be associative and commutative for the result to be well defined. The input tile is placed in shared memory automatically, so it counts against the block’s shared-memory budget (see Shared Memory Budget) even when the input uses register storage.

Parameters:

op – Callable taking two elements of the tile’s dtype and returning one value of the same type. May be a Warp built-in such as add() or a user-defined function decorated with @wp.func.
a – The input tile. Must have a scalar element type, unlike the whole-tile form, which also accepts vector and matrix elements. The operator, or one of its overloads, must accept the tile’s element type.
axis – The axis to reduce over. Must be a compile-time constant in [0, len(shape)).

Returns:

A register tile with the same dtype as a. For an input with two or more dimensions, the result has the shape of a with axis removed. For a one-dimensional input, it has shape (1,).

On a partial CPU block, a slice with no active values returns the operation’s identity for wp.add, wp.mul, wp.min, and wp.max. Other operators have no declared identity, so an empty slice triggers an assertion instead of returning an arbitrary value. See CPU Tile Semantics for definitions of partial CPU blocks and active lanes.

Example

@wp.kernel
def row_products(x: wp.array2d[float], p: wp.array[float]):
    t = wp.tile_load(x, shape=(2, 3))
    wp.tile_store(p, wp.tile_reduce(wp.mul, t, axis=1))

x = wp.array([[1.0, 2.0, 3.0], [10.0, 20.0, 30.0]], dtype=float)
p = wp.zeros(2, dtype=float)

wp.launch_tiled(row_products, dim=1, inputs=[x], outputs=[p], block_dim=4)

print(p.numpy())

[   6. 6000.]