warp.sum() Essentials - Warp-Level Dot Product

Implement the dot product we saw in puzzle 12 using Mojo’s warp operations to replace complex shared memory patterns with simple function calls. Each warp lane will process one element and use warp.sum() to combine results automatically, demonstrating how warp programming transforms GPU synchronization.

Key insight: The warp.sum() operation leverages SIMT execution to replace shared memory + barriers + tree reduction with a log2(WARP_SIZE)-step shuffle reduction behind one function call.

Key concepts

In this puzzle, you’ll learn:

  • Warp-level reductions with warp.sum()
  • SIMT execution model and lane synchronization
  • Cross-architecture compatibility with WARP_SIZE
  • Performance transformation from complex to simple patterns
  • Lane ID management and conditional writes

The mathematical operation is a dot product (inner product): \[\Large \text{output}[0] = \sum_{i=0}^{N-1} a[i] \times b[i]\]

But the implementation teaches fundamental patterns for all warp-level GPU programming in Mojo.

Configuration

  • Vector size: SIZE = WARP_SIZE (32 or 64 depending on GPU architecture)
  • Data type: DType.float32
  • Block configuration: (WARP_SIZE, 1) threads per block
  • Grid configuration: (1, 1) blocks per grid
  • Layout: row_major[SIZE]() (1D row-major)

Scope: This puzzle works within a single warp (SIZE = WARP_SIZE). The reduction happens across lanes of one warp via warp.sum(); there is no cross-warp or cross-block reduction here.

The traditional complexity (from Puzzle 12)

Recall the complex approach from solutions/p12/p12.mojo that required shared memory, barriers, and tree reduction:

def traditional_dot_product_p12_style[
    InLayoutT: TensorLayout, OutLayoutT: TensorLayout, size: Int
](
    output: TileTensor[mut=True, dtype, OutLayoutT, MutAnyOrigin],
    a: TileTensor[mut=False, dtype, InLayoutT, MutAnyOrigin],
    b: TileTensor[mut=False, dtype, InLayoutT, MutAnyOrigin],
):
    """
    This is the complex approach from p12_layout_tensor.mojo - kept for comparison.
    """
    var a_lt = a.to_layout_tensor()
    var b_lt = b.to_layout_tensor()
    var out_lt = output.to_layout_tensor()
    var shared = stack_allocation[
        dtype=dtype, address_space=AddressSpace.SHARED
    ](row_major[WARP_SIZE]())
    var global_i = block_dim.x * block_idx.x + thread_idx.x
    var local_i = thread_idx.x

    if global_i < size:
        shared[local_i] = rebind[Scalar[dtype]](a_lt[global_i]) * rebind[
            Scalar[dtype]
        ](b_lt[global_i])
    else:
        shared[local_i] = 0.0

    barrier()

    var stride = WARP_SIZE // 2
    while stride > 0:
        if local_i < stride:
            shared[local_i] += shared[local_i + stride]
        barrier()
        stride //= 2

    if local_i == 0:
        out_lt.store[1](Index(global_i // WARP_SIZE), shared[0])

What makes this complex:

  • Shared memory allocation: Manual memory management within blocks
  • Explicit barriers: barrier() calls to synchronize threads
  • Tree reduction: Complex loop with stride-based indexing
  • Conditional writes: Only thread 0 writes the final result

This works, but it’s verbose, error-prone, and requires deep understanding of GPU synchronization.

Note: This is intentionally a different approach from the Puzzle 12 solution. Puzzle 12 uses shared memory, barrier(), and a tree reduction; this puzzle deliberately replaces all of that with a single warp.sum(). The code below won’t match the P12 solution line-for-line — that contrast is the point.

Test the traditional approach:

pixi run p24 --traditional

pixi run -e amd p24 --traditional

pixi run -e apple p24 --traditional

uv run poe p24 --traditional

Code to complete

1. Simple warp kernel approach

Transform the complex traditional approach into a simple warp kernel using warp_sum():

def simple_warp_dot_product[
    InLayoutT: TensorLayout, OutLayoutT: TensorLayout, size: Int
](
    output: TileTensor[mut=True, dtype, OutLayoutT, MutAnyOrigin],
    a: TileTensor[mut=False, dtype, InLayoutT, MutAnyOrigin],
    b: TileTensor[mut=False, dtype, InLayoutT, MutAnyOrigin],
):
    var a_lt = a.to_layout_tensor()
    var b_lt = b.to_layout_tensor()
    var out_lt = output.to_layout_tensor()
    var global_i = block_dim.x * block_idx.x + thread_idx.x
    # FILL IN (6 lines at most)

View full file: problems/p24/p24.mojo

Tips

1. Understanding the simple warp kernel structure

You need to complete the simple_warp_dot_product function with 6 lines or fewer:

def simple_warp_dot_product[...](output, a, b):
    var global_i = block_dim.x * block_idx.x + thread_idx.x
    # FILL IN (6 lines at most)

Pattern to follow:

  1. Compute partial product for this thread’s element
  2. Use warp_sum() to combine across all warp lanes
  3. Lane 0 writes the final result

2. Computing partial products

var partial_product: Scalar[dtype] = 0
if global_i < size:
    partial_product = rebind[Scalar[dtype]](a_lt[global_i]) * rebind[
        Scalar[dtype]
    ](b_lt[global_i])

Why rebind? Indexing a LayoutTensor yields a SIMD value, so each element is narrowed to Scalar[dtype] with rebind before the multiply. Note the receivers are a_lt and b_lt—the LayoutTensor handles obtained from to_layout_tensor() inside the kernel, not the TileTensor parameters.

Bounds checking: Essential because not all threads may have valid data to process.

3. Warp reduction magic

var total = warp_sum(partial_product)

What warp_sum() does:

  • Takes each lane’s partial_product value
  • Sums them across all lanes in the warp with log2(WARP_SIZE) shuffle steps
  • Returns the same total to all lanes (not just lane 0)
  • Requires zero explicit synchronization (SIMT handles it)

4. Writing the result

if lane_id() == 0:
    out_lt.store[1](Index(global_i // WARP_SIZE), total)

Why only lane 0? All lanes have the same total value after warp_sum(), but we only want to write once to avoid race conditions.

Why not write to output[0]? Flexibility, function can be used in cases where there is more than one warp. i.e. The result from each warp is written to the unique location global_i // WARP_SIZE.

lane_id(): Returns 0 to WARP_SIZE - 1 (0-31 on NVIDIA, AMD RDNA and Apple; 0-63 on AMD CDNA) - identifies which lane within the warp.

Test the simple warp kernel:

uv run poe p24 --kernel

pixi run p24 --kernel

Expected output when solved:

SIZE: 32
WARP_SIZE: 32
SIMD_WIDTH: 8
=== RESULT ===
actual: HostBuffer([10416.0])
expected: HostBuffer([10416.0])
Puzzle 24 complete ✅

Solution

def simple_warp_dot_product[
    InLayoutT: TensorLayout, OutLayoutT: TensorLayout, size: Int
](
    output: TileTensor[mut=True, dtype, OutLayoutT, MutAnyOrigin],
    a: TileTensor[mut=False, dtype, InLayoutT, MutAnyOrigin],
    b: TileTensor[mut=False, dtype, InLayoutT, MutAnyOrigin],
):
    var a_lt = a.to_layout_tensor()
    var b_lt = b.to_layout_tensor()
    var out_lt = output.to_layout_tensor()
    var global_i = block_dim.x * block_idx.x + thread_idx.x

    # Each thread computes one partial product using vectorized approach as values in Mojo are SIMD based
    var partial_product: Scalar[dtype] = 0
    if global_i < size:
        partial_product = rebind[Scalar[dtype]](a_lt[global_i]) * rebind[
            Scalar[dtype]
        ](b_lt[global_i])

    # warp_sum() replaces all the shared memory + barriers + tree reduction
    var total = warp_sum(partial_product)

    # Only lane 0 writes the result (all lanes have the same total)
    if lane_id() == 0:
        out_lt.store[1](Index(global_i // WARP_SIZE), total)

The simple warp kernel demonstrates the fundamental transformation from explicit shared-memory synchronization to a single warp primitive:

What disappeared from the traditional approach:

  • 15+ lines → 6 lines: Dramatic code reduction
  • Shared memory allocation: Zero memory management required
  • 1 + log2(WARP_SIZE) barrier() calls: Zero explicit synchronization
  • Complex tree reduction: Single function call
  • Stride-based indexing: Eliminated entirely

SIMT execution model:

Warp lanes (SIMT execution):
Lane 0: partial_product = a[0] * b[0]    = 0.0
Lane 1: partial_product = a[1] * b[1]    = 1.0
Lane 2: partial_product = a[2] * b[2]    = 4.0
...
Lane 31: partial_product = a[31] * b[31] = 961.0

warp_sum() shuffle reduction:
All lanes → 0.0 + 1.0 + 4.0 + ... + 961.0 = 10416.0
All lanes receive → total = 10416.0 (broadcast result)

Why this works without barriers:

  1. SIMT execution: All lanes execute each instruction simultaneously
  2. Implicit synchronization: When warp_sum() begins, all lanes have computed their partial_product
  3. Register-to-register exchange: Lanes trade values through shuffle instructions instead of shared memory, so no barrier is needed. The reduction itself is a log2(WARP_SIZE)-step loop over those shuffles, not a single hardware reduction unit
  4. Broadcast result: All lanes receive the same total value

2. Functional approach

Now implement the same warp dot product using Mojo’s functional programming patterns:

def functional_warp_dot_product[
    InLayoutT: TensorLayout,
    OutLayoutT: TensorLayout,
    //,
    dtype: DType,
    simd_width: Int,
    rank: Int,
    size: Int,
](
    output: TileTensor[mut=True, dtype, OutLayoutT, MutAnyOrigin],
    a: TileTensor[mut=False, dtype, InLayoutT, MutAnyOrigin],
    b: TileTensor[mut=False, dtype, InLayoutT, MutAnyOrigin],
    ctx: DeviceContext,
) raises:
    @parameter
    @always_inline
    def compute_dot_product[
        simd_width: Int, alignment: Int = align_of[dtype]()
    ](indices: Coord) capturing -> None:
        var idx = Int(indices[0].value())
        # Convert inside GPU kernel to avoid host-captured LayoutTensor issues
        var a_lt = a.to_layout_tensor()
        var b_lt = b.to_layout_tensor()
        var out_lt = output.to_layout_tensor()
        # FILL IN (10 lines at most)

    # Launch exactly size == WARP_SIZE threads (one warp) to process all elements
    elementwise[compute_dot_product, 1, target="gpu"](size, ctx)
Tips

1. Understanding the functional approach structure

You need to complete the compute_dot_product function with 10 lines or fewer:

@parameter
@always_inline
def compute_dot_product[
    simd_width: Int, alignment: Int = align_of[dtype]()
](indices: Coord) capturing -> None:
    var idx = Int(indices[0].value())
    # FILL IN (10 lines at most)

Functional pattern differences:

  • Uses elementwise to launch exactly WARP_SIZE threads
  • Each thread processes one element based on idx
  • Same warp operations, different launch mechanism

2. Computing partial products

var partial_product: Scalar[dtype] = 0.0
if idx < size:
    var a_val = a_lt.load[1](Index(idx))
    var b_val = b_lt.load[1](Index(idx))
    partial_product = a_val * b_val
else:
    partial_product = 0.0

Loading pattern: a_lt.load[1](Index(idx)) loads exactly 1 element at position idx (not SIMD vectorized). The tensor is 1-D, so the index is a single Index(idx).

Bounds handling: Set partial_product = 0.0 for out-of-bounds threads so they don’t contribute to the sum.

3. Warp operations and storing

var total = warp_sum(partial_product)

if lane_id() == 0:
    out_lt.store[1](Index(idx // WARP_SIZE), total)

Storage pattern: out_lt.store[1](Index(idx // WARP_SIZE), total) stores 1 element at position idx // WARP_SIZE in the 1-D output tensor.

Same warp logic: warp_sum() and lane 0 writing work identically in functional approach.

4. Available functions from imports

from std.gpu import lane_id
from std.gpu.primitives.warp import sum as warp_sum, WARP_SIZE

# Inside your function:
var my_lane = lane_id()           # 0 to WARP_SIZE-1
var total = warp_sum(my_value)    # Shuffle-based reduction
var warp_size = WARP_SIZE         # 32 on NVIDIA, AMD RDNA, Apple; 64 on CDNA

Test the functional approach:

uv run poe p24 --functional

pixi run p24 --functional

Expected output when solved:

SIZE: 32
WARP_SIZE: 32
SIMD_WIDTH: 8
=== RESULT ===
actual: HostBuffer([10416.0])
expected: HostBuffer([10416.0])
Puzzle 24 complete ✅

Solution

def functional_warp_dot_product[
    InLayoutT: TensorLayout,
    OutLayoutT: TensorLayout,
    //,
    dtype: DType,
    simd_width: Int,
    rank: Int,
    size: Int,
](
    output: TileTensor[mut=True, dtype, OutLayoutT, MutAnyOrigin],
    a: TileTensor[mut=False, dtype, InLayoutT, MutAnyOrigin],
    b: TileTensor[mut=False, dtype, InLayoutT, MutAnyOrigin],
    ctx: DeviceContext,
) raises:
    @parameter
    @always_inline
    def compute_dot_product[
        simd_width: Int, alignment: Int = align_of[dtype]()
    ](indices: Coord) capturing -> None:
        var idx = Int(indices[0].value())
        # Convert inside GPU kernel to avoid host-captured LayoutTensor issues
        var a_lt = a.to_layout_tensor()
        var b_lt = b.to_layout_tensor()
        var out_lt = output.to_layout_tensor()

        # Each thread computes one partial product
        var partial_product: Scalar[dtype] = 0.0
        if idx < size:
            var a_val = a_lt.load[1](Index(idx))
            var b_val = b_lt.load[1](Index(idx))
            partial_product = a_val * b_val
        else:
            partial_product = 0.0

        # Warp magic - combines all WARP_SIZE partial products!
        var total = warp_sum(partial_product)

        # Only lane 0 writes the result (all lanes have the same total)
        if lane_id() == 0:
            out_lt.store[1](Index(idx // WARP_SIZE), total)

    # Launch exactly size == WARP_SIZE threads (one warp) to process all elements
    elementwise[compute_dot_product, 1, target="gpu"](size, ctx)

The functional warp approach showcases modern Mojo programming patterns with warp operations:

Functional approach characteristics:

elementwise[compute_dot_product, 1, target="gpu"](size, ctx)

Benefits:

  • Type safety: Compile-time tensor layout checking
  • Composability: Easy integration with other functional operations
  • Modern patterns: Leverages Mojo’s functional programming features
  • Automatic optimization: Compiler can apply high-level optimizations

Key differences from kernel approach:

  • Launch mechanism: Uses elementwise instead of enqueue_function
  • Memory access: Uses .load[1]() and .store[1]() patterns
  • Integration: Seamlessly works with other functional operations

Same warp benefits:

  • Zero synchronization: warp_sum() works identically
  • Same reduction: warp_sum() expands to the same shuffle sequence in both approaches; only the launch mechanism differs
  • Cross-architecture: WARP_SIZE adapts automatically

Performance comparison with benchmarks

Run comprehensive benchmarks to see how warp operations scale:

uv run poe p24 --benchmark

pixi run p24 --benchmark

Here’s example output from a complete benchmark run:

SIZE: 32
WARP_SIZE: 32
SIMD_WIDTH: 8
--------------------------------------------------------------------------------
Testing SIZE=1 x WARP_SIZE, BLOCKS=1
Running traditional_1x
Running simple_warp_1x
Running functional_warp_1x
--------------------------------------------------------------------------------
Testing SIZE=4 x WARP_SIZE, BLOCKS=4
Running traditional_4x
Running simple_warp_4x
Running functional_warp_4x
--------------------------------------------------------------------------------
Testing SIZE=32 x WARP_SIZE, BLOCKS=32
Running traditional_32x
Running simple_warp_32x
Running functional_warp_32x
--------------------------------------------------------------------------------
Testing SIZE=256 x WARP_SIZE, BLOCKS=256
Running traditional_256x
Running simple_warp_256x
Running functional_warp_256x
--------------------------------------------------------------------------------
Testing SIZE=2048 x WARP_SIZE, BLOCKS=2048
Running traditional_2048x
Running simple_warp_2048x
Running functional_warp_2048x
--------------------------------------------------------------------------------
Testing SIZE=16384 x WARP_SIZE, BLOCKS=16384 (Large Scale)
Running traditional_16384x
Running simple_warp_16384x
Running functional_warp_16384x
--------------------------------------------------------------------------------
Testing SIZE=65536 x WARP_SIZE, BLOCKS=65536 (Massive Scale)
Running traditional_65536x
Running simple_warp_65536x
Running functional_warp_65536x
| name                   | met (ms)              | iters |
| ---------------------- | --------------------- | ----- |
| traditional_1x         | 0.00460128            | 100   |
| simple_warp_1x         | 0.00574047            | 100   |
| functional_warp_1x     | 0.00484192            | 100   |
| traditional_4x         | 0.00492671            | 100   |
| simple_warp_4x         | 0.00485247            | 100   |
| functional_warp_4x     | 0.00587679            | 100   |
| traditional_32x        | 0.0062406399999999996 | 100   |
| simple_warp_32x        | 0.0054918400000000004 | 100   |
| functional_warp_32x    | 0.00552447            | 100   |
| traditional_256x       | 0.0050614300000000004 | 100   |
| simple_warp_256x       | 0.00488768            | 100   |
| functional_warp_256x   | 0.00461472            | 100   |
| traditional_2048x      | 0.01120031            | 100   |
| simple_warp_2048x      | 0.00884383            | 100   |
| functional_warp_2048x  | 0.007038720000000001  | 100   |
| traditional_16384x     | 0.038533750000000005  | 100   |
| simple_warp_16384x     | 0.0323264             | 100   |
| functional_warp_16384x | 0.01674271            | 100   |
| traditional_65536x     | 0.19784991999999998   | 100   |
| simple_warp_65536x     | 0.12870176            | 100   |
| functional_warp_65536x | 0.048680310000000004  | 100   |

Benchmarks completed!

WARP OPERATIONS PERFORMANCE ANALYSIS:
   GPU Architecture: NVIDIA (WARP_SIZE=32) vs AMD (WARP_SIZE=64)
   - 1,...,256 x WARP_SIZE: Grid size too small to benchmark
   - 2048 x WARP_SIZE: Warp primitive benefits emerge
   - 16384 x WARP_SIZE: Large scale (512K-1M elements)
   - 65536 x WARP_SIZE: Massive scale (2M-4M elements)

   Expected Results at Large Scales:
   • Traditional: Slower due to more barrier overhead
   • Warp operations: Faster, scale better with problem size
   • Memory bandwidth becomes the limiting factor

Performance insights from this example:

  • Small scales (1x-256x): All three land within a few microseconds of each other and the ordering flips between runs - the grid is too small for the measurement to separate them, which is why the analysis text says so
  • Medium scale (2048x): The warp approaches start to pull ahead of the traditional tree reduction
  • Large scales (16K-65K): The gap widens rather than closing - in this run the functional warp approach finishes about 4x faster than the traditional one at 65536x
  • Variability: Performance depends heavily on specific GPU architecture and memory subsystem

Note: Your results will vary significantly depending on your hardware (GPU model, memory bandwidth, WARP_SIZE). The key insight is observing the relative performance trends rather than absolute timings.

Next steps

Once you’ve learned warp sum operations, you’re ready for:

  • When to Use Warp Programming: Strategic decision framework for warp vs traditional approaches
  • Advanced warp operations: shuffle_idx(), shuffle_down(), prefix_sum() for complex communication patterns
  • Multi-warp algorithms: Combining warp operations with block-level synchronization
  • Part VIII: Block-Level Programming: Scaling these patterns from one warp to a whole block

💡 Key Takeaway: Warp operations transform GPU programming by replacing complex synchronization patterns with lane-to-lane shuffles, demonstrating how understanding the execution model enables dramatic simplification without sacrificing performance.