Add wp.tile_dot() for fused element-wise multiply and reduce

Description:

Computing a dot product between two tiles currently requires three separate operations: wp.tile_map(wp.mul, a, b) to create an element-wise product tile, wp.tile_sum() to reduce it, and wp.tile_extract() to get the scalar result. This creates intermediate tiles (both register and shared) that consume registers and shared memory unnecessarily.

A dedicated wp.tile_dot(a, b) primitive would:

Return a scalar value directly (not a tile), broadcast to all threads
Create no intermediate tiles — fuse the multiply, local reduction, and warp reduction into one operation
Reduce register pressure and shared memory usage in tight inner loops

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

Closed by commit 5bef033.

scalar = wp.tile_dot(a, b)  # full contraction, broadcast to all threads

a, b: tiles with the same shape and dtype.
Returns the tile's scalar type (so a tile of vec3f returns float), matching PyTorch / JAX / Triton sum() semantics.
Internally performs a full tensordot per element pair, then a block-wide sum, with no intermediate tiles or tile_extract round trip. Equivalent to wp.tile_extract(wp.tile_sum(wp.tile_map(wp.tensordot, a, b)), 0).

@wp.kernel
def compute():
    a = wp.tile_ones(dtype=float, shape=64)
    b = wp.tile_ones(dtype=float, shape=64) * 2.0
    d = wp.tile_dot(a, b)
    print(d)

wp.launch_tiled(compute, dim=[1], inputs=[], block_dim=64)
# 128.0

Follow-up: an adjoint bug was flagged during MR review but deferred so the primitive could land. Tracked in sub-issue #1374.

Navigation Menu

Repository navigation

Add `wp.tile_dot()` for fused element-wise multiply and reduce #1364

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Add wp.tile_dot() for fused element-wise multiply and reduce #1364

Description

Activity

shi-eric commented on Apr 19, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

Add `wp.tile_dot()` for fused element-wise multiply and reduce #1364