warp.tile_dot#

warp.tile_dot( a: Tile[Any, tuple[int, ...]], b: Tile[Any, tuple[int, ...]], ) → Tile[Any, tuple[Literal[1]]]#

Kernel: true
Python: false
Differentiable: true

Compute the dot product of two tiles.

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

Contracts corresponding elements and sums the results over the whole tile. For scalar tiles this is the standard dot product; for vector tiles each pair is contracted with dot(); for matrix tiles it is the Frobenius inner product.

Mathematically equivalent to wp.tile_sum(a * b) for scalar tiles, wp.tile_sum(wp.tile_map(wp.dot, a, b)) for vector tiles, and wp.tile_sum(wp.tile_map(wp.ddot, a, b)) for matrix tiles, but without the intermediate tile those forms would materialize. The summation order is not guaranteed to match, so floating-point results may differ in the last bits.

Parameters:

a – First tile operand.
b – Second tile operand with the same shape and dtype as a. The element type must be scalar, vector, or matrix.

Returns:

A shape-(1,) register tile holding the result, with the scalar type of the operands’ elements. Read the value with [0].

Example

@wp.kernel
def dot_product(a: wp.array[float], b: wp.array[float], out: wp.array[float]):
    ta = wp.tile_load(a, shape=4)
    tb = wp.tile_load(b, shape=4)
    wp.tile_store(out, wp.tile_dot(ta, tb))

a = wp.array([1.0, 2.0, 3.0, 4.0], dtype=float)
b = wp.array([10.0, 20.0, 30.0, 40.0], dtype=float)
out = wp.zeros(1, dtype=float)

wp.launch_tiled(dot_product, dim=1, inputs=[a, b], outputs=[out], block_dim=2)

print(out.numpy())

[300.]