Add wp.tile_axpy() for fused scalar-multiply-and-accumulate on tiles

Description:

There is currently no built-in way to perform a fused dest += src * alpha operation on tiles. Users must decompose this into tile_map(mul, src, alpha_tile) followed by element-wise addition, which creates intermediate register tiles and may prevent the compiler from generating optimal fused multiply-add (FMA) instructions.

A dedicated wp.tile_axpy(dest, src, alpha) primitive would:

Perform dest[i] += src[i] * alpha in a single pass with no intermediate tiles
Work with mixed storage (e.g., shared dest + register src)
Enable the compiler to use FMA instructions throughout
Match the BLAS axpy pattern commonly needed in iterative solvers

Add wp.tile_axpy() and wp.tile_dot() fused tile primitives (NVIDIAGH-…

Closed by commit 5bef033.

wp.tile_axpy(alpha, src, dest)  # dest += src * alpha (in place)

alpha: scalar, must match the tile's underlying scalar type (e.g. float for a tile of vec3f).
src, dest: tiles with the same shape and dtype.
Fused multiply-add, no intermediate scaled tile.
For shared dest tiles with non-unique layouts (broadcast bias tiles) the accumulate is done atomically to avoid write races.
Differentiable (forward and reverse implemented).

@wp.kernel
def compute():
    dest = wp.tile_ones(dtype=float, shape=4) * 2.0
    src  = wp.tile_ones(dtype=float, shape=4) * 3.0
    wp.tile_axpy(5.0, src, dest)
    print(dest)

wp.launch_tiled(compute, dim=[1], inputs=[], block_dim=64)
# [17 17 17 17] = tile(shape=(4), storage=register)

Navigation Menu

Repository navigation

Add `wp.tile_axpy()` for fused scalar-multiply-and-accumulate on tiles #1363

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Add wp.tile_axpy() for fused scalar-multiply-and-accumulate on tiles #1363

Description

Activity

shi-eric commented on Apr 19, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

Add `wp.tile_axpy()` for fused scalar-multiply-and-accumulate on tiles #1363