warp.tile_map#

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

Kernel

Differentiable

Apply a unary function onto the tile.

This function cooperatively applies a unary function to each element of the tile using all threads in the block.

param op:

A callable function that accepts one argument and returns one argument, may be a user function or builtin

param a:

The input tile, the operator (or one of its overloads) must be able to accept the tile’s data type

returns:

A tile with the same dimensions as the input tile. Its datatype is specified by the return type of op

Example:
@wp.kernel
def compute():

    t = wp.tile_arange(0.0, 1.0, 0.1, dtype=float)
    s = wp.tile_map(wp.sin, t)

    print(s)

wp.launch_tiled(compute, dim=[1], inputs=[], block_dim=16)
Prints:
[0 0.0998334 0.198669 0.29552 0.389418 0.479426 0.564642 0.644218 0.717356 0.783327] = tile(shape=(10), storage=register)

warp.tile_map( op: Callable, a: Tile[Any, tuple[int, ...]], b: Tile[Any, tuple[int, ...]], ) → Tile[Any, tuple[int, ...]]

Kernel

Differentiable

Apply a binary function onto the tile.

This function cooperatively applies a binary function to each element of the tiles using all threads in the block. Both input tiles must have the same dimensions, and if using a builtin op, the same datatypes.

param op:

A callable function that accepts two arguments and returns one argument, all of the same type, may be a user function or builtin

param a:

The first input tile, the operator (or one of its overloads) must be able to accept the tile’s dtype

param b:

The second input tile, the operator (or one of its overloads) must be able to accept the tile’s dtype

returns:

A tile with the same dimensions as the input tiles. Its datatype is specified by the return type of op

Example:
@wp.kernel
def compute():

    a = wp.tile_arange(0.0, 1.0, 0.1, dtype=float)
    b = wp.tile_ones(shape=10, dtype=float)

    s = wp.tile_map(wp.add, a, b)

    print(s)

wp.launch_tiled(compute, dim=[1], inputs=[], block_dim=16)
Prints:
[1 1.1 1.2 1.3 1.4 1.5 1.6 1.7 1.8 1.9] = tile(shape=(10), storage=register)

warp.tile_map( op: Callable, *args: Tile[Scalar, tuple[int, ...]], ) → Tile[Scalar, tuple[int, ...]]

Kernel

Differentiable

Apply a user-defined function to multiple tiles element-wise.

This function cooperatively applies a user-defined function to corresponding elements of three or more tiles using all threads in the block. All input tiles must have the same dimensions. The operator must accept the same number of arguments as tiles provided.

param op:

A callable function that accepts N arguments and returns one value, must be a user function

param args:

Three or more input tiles with matching dimensions. The operator (or one of its overloads) must be able to accept the tiles’ dtypes

returns:

A tile with the same dimensions as the input tiles. Its datatype is specified by the return type of op

Example:
@wp.func
def weighted_sum(a: float, b: float, c: float):
    return 0.5 * a + 0.3 * b + 0.2 * c

@wp.kernel
def compute():

    a = wp.tile_arange(0.0, 1.0, 0.1, dtype=float)
    b = wp.tile_ones(shape=10, dtype=float)
    c = wp.tile_arange(1.0, 2.0, 0.1, dtype=float)

    s = wp.tile_map(weighted_sum, a, b, c)

    print(s)

wp.launch_tiled(compute, dim=[1], inputs=[], block_dim=16)
Prints:
[0.5 0.57 0.64 0.71 0.78 0.85 0.92 0.99 1.06 1.13] = tile(shape=(10), storage=register)