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 a tile.

Reduce across all elements using the provided operator.

Parameters:

op – A callable function that accepts two arguments and returns one argument, may be a user function or builtin
a – The input tile, the operator (or one of its overloads) must be able to accept the tile’s data type

Returns:

A single-element tile with the same data type as the input tile.

Example

@wp.kernel
def factorial():

    t = wp.tile_arange(1, 10, dtype=int)
    s = wp.tile_reduce(wp.mul, t)

    print(s)

wp.launch_tiled(factorial, dim=[1], inputs=[], block_dim=16)

[362880] = tile(shape=(1), storage=register)

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 across a tile.

Reduce across a tile axis using the provided operator.

Parameters:

op – A callable function that accepts two arguments and returns one argument, may be a user function or builtin
a – The input tile, the operator (or one of its overloads) must be able to accept the tile’s data type. Must reside in shared memory.
axis – The tile axis to perform the reduction across. Must be a compile-time constant.

Returns:

A tile with the same shape as the input tile less the axis dimension and the same data type as the input tile.

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

TILE_M = wp.constant(4)
TILE_N = wp.constant(2)

@wp.kernel
def compute(x: wp.array2d[float], y: wp.array[float]):

    a = wp.tile_load(x, shape=(TILE_M, TILE_N))
    b = wp.tile_reduce(wp.add, a, axis=1)
    wp.tile_store(y, b)

arr = np.arange(TILE_M * TILE_N).reshape(TILE_M, TILE_N)

x = wp.array(arr, dtype=float)
y = wp.zeros(TILE_M, dtype=float)

wp.launch_tiled(compute, dim=[1], inputs=[x], outputs=[y], block_dim=32)

print(x.numpy())
print(y.numpy())

[[0. 1.]
 [2. 3.]
 [4. 5.]
 [6. 7.]]
[ 1.  5.  9. 13.]