Vectorization - Fine-Grained SIMD Control

Overview

This puzzle explores advanced vectorization techniques using manual vectorization and vectorize that give you precise control over SIMD operations within GPU kernels. You’ll implement two different approaches to vectorized computation:

  1. Manual vectorization: Direct SIMD control with explicit index calculations
  2. Mojo’s vectorize function: High-level vectorization with automatic remainder handling

Both approaches build on tiling concepts but with different trade-offs between control, convenience, and performance optimization.

Key insight: Different vectorization strategies suit different performance requirements and complexity levels.

Key concepts

In this puzzle, you’ll learn:

  • Manual SIMD operations with explicit index management
  • Mojo’s vectorize function for automatic chunk and remainder handling
  • Chunk-based memory organization for optimal SIMD alignment
  • Bounds checking strategies for edge cases
  • Performance trade-offs between manual control and generated loops

The same mathematical operation as before: \[\Large \text{output}[i] = a[i] + b[i]\]

But with sophisticated vectorization strategies for maximum performance.

Configuration

  • Vector size: SIZE = 1024
  • Tile size: TILE_SIZE = 32
  • Data type: DType.float32
  • SIMD width: GPU-dependent
  • Layout: row_major[SIZE]() (1D row-major)

Scope: Both approaches operate within a single tile at a time; bounds checking is per-tile and there is no cross-tile or cross-block communication. The focus is SIMD control inside a tile, not coordination across them.

1. Manual vectorization approach

Code to complete

def manual_vectorized_tiled_elementwise_add[
    LayoutT: TensorLayout,
    dtype: DType,
    simd_width: Int,
    num_threads_per_tile: Int,
    rank: Int,
    size: Int,
    tile_size: Int,
](
    output: TileTensor[mut=True, dtype, LayoutT, MutAnyOrigin],
    a: TileTensor[mut=False, dtype, LayoutT, MutAnyOrigin],
    b: TileTensor[mut=False, dtype, LayoutT, MutAnyOrigin],
    ctx: DeviceContext,
) raises:
    # Each tile contains tile_size groups of simd_width elements
    comptime chunk_size = tile_size * simd_width

    @parameter
    @always_inline
    def process_manual_vectorized_tiles[
        num_threads_per_tile: Int, alignment: Int = align_of[dtype]()
    ](indices: Coord) capturing -> None:
        var tile_id = 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 (7 lines at most)

    # Number of tiles needed: each tile processes chunk_size elements
    var num_tiles = (size + chunk_size - 1) // chunk_size
    elementwise[
        process_manual_vectorized_tiles, num_threads_per_tile, target="gpu"
    ](num_tiles, ctx)

View full file: problems/p23/p23.mojo

Tips

1. Understanding chunk organization

comptime chunk_size = tile_size * simd_width  # 32 * 4 = 128 elements per chunk

Each tile now contains multiple SIMD groups, not just sequential elements.

2. Global index calculation

var global_start = tile_id * chunk_size + i * simd_width

This calculates the exact global position for each SIMD vector within the chunk.

3. Direct tensor access

var a_vec = a_lt.aligned_load[width=simd_width](Index(global_start))     # Load from global tensor
out_lt.store[simd_width](Index(global_start), ret)  # Store to global tensor

Note: Access the whole-tensor LayoutTensor handles from to_layout_tensor(), not the tile views.

4. Key characteristics

  • More control, more complexity, global tensor access
  • Perfect SIMD alignment with hardware
  • No bounds check at all: size must divide evenly into chunk_size

Running manual vectorization

pixi run p23 --manual-vectorized

pixi run -e amd p23 --manual-vectorized

pixi run -e apple p23 --manual-vectorized

uv run poe p23 --manual-vectorized

Your output will look like this when not yet solved:

SIZE: 1024
simd_width: 4
tile size: 32
out: HostBuffer([0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0])
expected: HostBuffer([1.0, 5.0, 9.0, ..., 4085.0, 4089.0, 4093.0])

Manual vectorization solution

def manual_vectorized_tiled_elementwise_add[
    LayoutT: TensorLayout,
    dtype: DType,
    simd_width: Int,
    num_threads_per_tile: Int,
    rank: Int,
    size: Int,
    tile_size: Int,
](
    output: TileTensor[mut=True, dtype, LayoutT, MutAnyOrigin],
    a: TileTensor[mut=False, dtype, LayoutT, MutAnyOrigin],
    b: TileTensor[mut=False, dtype, LayoutT, MutAnyOrigin],
    ctx: DeviceContext,
) raises:
    # Each tile contains tile_size groups of simd_width elements
    comptime chunk_size = tile_size * simd_width

    @parameter
    @always_inline
    def process_manual_vectorized_tiles[
        num_threads_per_tile: Int, alignment: Int = align_of[dtype]()
    ](indices: Coord) capturing -> None:
        var tile_id = 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()

        comptime for i in range(tile_size):
            var global_start = tile_id * chunk_size + i * simd_width

            var a_vec = a_lt.aligned_load[width=simd_width](Index(global_start))
            var b_vec = b_lt.aligned_load[width=simd_width](Index(global_start))
            var ret = a_vec + b_vec
            out_lt.store[simd_width](Index(global_start), ret)

    # Number of tiles needed: each tile processes chunk_size elements
    var num_tiles = (size + chunk_size - 1) // chunk_size
    elementwise[
        process_manual_vectorized_tiles, num_threads_per_tile, target="gpu"
    ](num_tiles, ctx)

Manual vectorization deep dive

Manual vectorization gives you direct control over SIMD operations with explicit index calculations:

  • Chunk-based organization: chunk_size = tile_size * simd_width
  • Global indexing: Direct calculation of memory positions
  • Manual bounds management: You handle edge cases explicitly

Architecture and memory layout:

comptime chunk_size = tile_size * simd_width  # 32 * 4 = 128

Chunk organization visualization (TILE_SIZE=32, SIMD_WIDTH=4):

Original array: [0, 1, 2, 3, ..., 1023]

Chunk 0 (thread 0): [0:128]    ← 128 elements = 32 SIMD groups of 4
Chunk 1 (thread 1): [128:256]  ← Next 128 elements
Chunk 2 (thread 2): [256:384]  ← Next 128 elements
...
Chunk 7 (thread 7): [896:1024] ← Final 128 elements

Processing within one chunk:

comptime for i in range(tile_size):  # i = 0, 1, 2, ..., 31
    var global_start = tile_id * chunk_size + i * simd_width
    # For tile_id=0: global_start = 0, 4, 8, 12, ..., 124
    # For tile_id=1: global_start = 128, 132, 136, 140, ..., 252

Performance characteristics:

  • Thread count: 8 threads (1024 ÷ 128 = 8)
  • Work per thread: 128 elements (32 SIMD operations of 4 elements each)
  • Memory pattern: Large chunks with perfect SIMD alignment
  • Overhead: Minimal - direct hardware mapping
  • Safety: No bounds check; correct only because size divides evenly into chunk_size

Key advantages:

  • Predictable indexing: Exact control over memory access patterns
  • Optimal alignment: SIMD operations perfectly aligned to hardware
  • No per-element guard: The chunk loop runs unconditionally
  • Hardware optimization: Direct mapping to GPU SIMD units

Key challenges:

  • Index complexity: Manual calculation of global positions
  • Bounds responsibility: A ragged tail is yours to handle, and this solution assumes there isn’t one
  • Debugging difficulty: More complex to verify correctness

2. Mojo vectorize approach

Code to complete

def vectorize_within_tiles_elementwise_add[
    LayoutT: TensorLayout,
    dtype: DType,
    simd_width: Int,
    num_threads_per_tile: Int,
    rank: Int,
    size: Int,
    tile_size: Int,
](
    output: TileTensor[mut=True, dtype, LayoutT, MutAnyOrigin],
    a: TileTensor[mut=False, dtype, LayoutT, MutAnyOrigin],
    b: TileTensor[mut=False, dtype, LayoutT, MutAnyOrigin],
    ctx: DeviceContext,
) raises:
    # Each tile contains tile_size elements (not SIMD groups)
    @parameter
    @always_inline
    def process_tile_with_vectorize[
        num_threads_per_tile: Int, alignment: Int = align_of[dtype]()
    ](indices: Coord) capturing -> None:
        var tile_id = Int(indices[0].value())
        var tile_start = tile_id * tile_size
        var tile_end = min(tile_start + tile_size, size)
        var actual_tile_size = tile_end - tile_start
        # 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 (9 lines at most)

    var num_tiles = (size + tile_size - 1) // tile_size
    elementwise[
        process_tile_with_vectorize, num_threads_per_tile, target="gpu"
    ](num_tiles, ctx)

View full file: problems/p23/p23.mojo

Tips

0. From scalar to vectorized

Start by writing the addition as a plain scalar loop over a tile, then convert it to vectorize. The transformation is mechanical: replace the per-element loop body with a SIMD load/add/store, and hand the loop to vectorize, which calls your body in width-sized steps and then calls it once per leftover element with width=1.

# Before: scalar loop over the tile (one element at a time)
for i in range(actual_tile_size):
    var global_idx = tile_start + i
    out_lt[global_idx] = a_lt[global_idx] + b_lt[global_idx]

# After: same logic, but the body operates on a SIMD vector of `width`
def vectorized_add[width: Int](i: Int) {imm tile_start, imm a_lt, imm b_lt, mut out_lt}:
    global_idx = tile_start + i
    if global_idx + width <= size:                       # bounds check
        var a_vec = a_lt.aligned_load[width](Index(global_idx))
        var b_vec = b_lt.aligned_load[width](Index(global_idx))
        out_lt.store[width](Index(global_idx), a_vec + b_vec)

vectorize[simd_width](actual_tile_size, vectorized_add)  # drives the loop + remainder

The remaining tips break this down piece by piece.

1. Tile boundary calculation

var tile_start = tile_id * tile_size
var tile_end = min(tile_start + tile_size, size)
var actual_tile_size = tile_end - tile_start

Handle cases where the last tile might be smaller than tile_size.

2. Vectorized function pattern

def vectorized_add[
  width: Int
](i: Int) {imm tile_start, imm a_lt, imm b_lt, mut out_lt}:
    var global_idx = tile_start + i
    if global_idx + width <= size:  # Bounds checking
        # SIMD operations here

The width parameter is automatically determined by the vectorize function.

3. Calling vectorize

vectorize[simd_width](actual_tile_size, vectorized_add)

This automatically handles the vectorization loop with the provided SIMD width.

4. Key characteristics

  • Automatic remainder handling, tile-based access
  • Takes explicit SIMD width parameter
  • Drives the loop and the remainder for you; the bounds check inside vectorized_add is still yours to write

Running Mojo vectorize

uv run poe p23 --vectorized

pixi run p23 --vectorized

Your output will look like this when not yet solved:

SIZE: 1024
simd_width: 4
tile size: 32
out: HostBuffer([0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0])
expected: HostBuffer([1.0, 5.0, 9.0, ..., 4085.0, 4089.0, 4093.0])

Mojo vectorize solution

def vectorize_within_tiles_elementwise_add[
    LayoutT: TensorLayout,
    dtype: DType,
    simd_width: Int,
    num_threads_per_tile: Int,
    rank: Int,
    size: Int,
    tile_size: Int,
](
    output: TileTensor[mut=True, dtype, LayoutT, MutAnyOrigin],
    a: TileTensor[mut=False, dtype, LayoutT, MutAnyOrigin],
    b: TileTensor[mut=False, dtype, LayoutT, MutAnyOrigin],
    ctx: DeviceContext,
) raises:
    # Each tile contains tile_size elements (not SIMD groups)
    @parameter
    @always_inline
    def process_tile_with_vectorize[
        num_threads_per_tile: Int, alignment: Int = align_of[dtype]()
    ](indices: Coord) capturing -> None:
        var tile_id = Int(indices[0].value())
        var tile_start = tile_id * tile_size
        var tile_end = min(tile_start + tile_size, size)
        var actual_tile_size = tile_end - tile_start
        # 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()

        def vectorized_add[
            width: Int
        ](i: Int) {imm tile_start, imm a_lt, imm b_lt, mut out_lt}:
            var global_idx = tile_start + i
            if global_idx + width <= size:
                var a_vec = a_lt.aligned_load[width](Index(global_idx))
                var b_vec = b_lt.aligned_load[width](Index(global_idx))
                var result = a_vec + b_vec
                out_lt.store[width](Index(global_idx), result)

        # Use vectorize within each tile
        vectorize[simd_width](actual_tile_size, vectorized_add)

    var num_tiles = (size + tile_size - 1) // tile_size
    elementwise[
        process_tile_with_vectorize, num_threads_per_tile, target="gpu"
    ](num_tiles, ctx)

Mojo vectorize deep dive

Mojo’s vectorize function drives the loop and its remainder for you:

  • Explicit SIMD width parameter: You provide the simd_width to use
  • Automatic remainder handling: Processes leftover elements automatically
  • Explicit bounds check: vectorize doesn’t validate indices, which is why vectorized_add keeps its own if global_idx + width <= size guard
  • Nested function pattern: Clean separation of vectorization logic

Tile-based organization:

var tile_start = tile_id * tile_size    # 0, 32, 64, 96, ...
var tile_end = min(tile_start + tile_size, size)
var actual_tile_size = tile_end - tile_start

Automatic vectorization mechanism:

def vectorized_add[
  width: Int
](i: Int) {imm tile_start, imm a_lt, imm b_lt, mut out_lt}:
    var global_idx = tile_start + i
    if global_idx + width <= size:
        # Automatic SIMD optimization

How vectorize works:

  • Automatic chunking: Divides actual_tile_size into chunks of your provided simd_width
  • Remainder handling: Calls your body once per leftover element with width=1
  • Loop management: Handles the vectorization loop automatically

Execution visualization (TILE_SIZE=32, SIMD_WIDTH=4):

Tile 0 processing:
  vectorize call 0: processes elements [0:4]   with SIMD_WIDTH=4
  vectorize call 1: processes elements [4:8]   with SIMD_WIDTH=4
  ...
  vectorize call 7: processes elements [28:32] with SIMD_WIDTH=4
  Total: 8 automatic SIMD operations

Performance characteristics:

  • Thread count: 32 threads (1024 ÷ 32 = 32)
  • Work per thread: 32 elements (automatic SIMD chunking)
  • Memory pattern: Smaller tiles with automatic vectorization
  • Overhead: Slight - the loop and remainder are generated for you
  • Safety: Edge cases handled by the generated remainder loop plus your own bounds guard

Performance comparison and best practices

When to use each approach

Choose manual vectorization when:

  • Maximum performance is critical
  • You have predictable, aligned data patterns
  • Expert-level control over memory access is needed
  • You can guarantee bounds safety manually
  • Hardware-specific optimization is required

Choose Mojo vectorize when:

  • Development speed and fewer hand-written loops are priorities
  • Working with irregular or dynamic data sizes
  • You want automatic remainder handling instead of manual edge case management
  • Hand-written chunk and remainder loops would be error-prone
  • You prefer cleaner vectorization patterns over manual loop management

Advanced optimization insights

Memory bandwidth utilization:

Manual:    8 threads × 32 SIMD ops = 256 total SIMD operations
Vectorize: 32 threads × 8 SIMD ops = 256 total SIMD operations

Both achieve similar total throughput but with different parallelism strategies.

Cache behavior:

  • Manual: Large chunks may exceed L1 cache, but perfect sequential access
  • Vectorize: Smaller tiles fit better in cache, with automatic remainder handling

Hardware mapping:

  • Manual: Direct control over warp utilization and SIMD unit mapping
  • Vectorize: Simplified vectorization with automatic loop and remainder management

Best practices summary

Manual vectorization best practices:

  • Always validate index calculations carefully
  • Use compile-time constants for chunk_size when possible
  • Profile memory access patterns for cache optimization
  • Consider alignment requirements for optimal SIMD performance

Mojo vectorize best practices:

  • Choose appropriate SIMD width for your data and hardware
  • Focus on algorithm clarity over micro-optimizations
  • Use nested parameter functions for clean vectorization logic
  • Let vectorize handle the remainder, but keep your own bounds guard in the body

Both approaches represent valid strategies in the GPU performance optimization toolkit, with manual vectorization offering maximum control and Mojo’s vectorize handling the loop and its remainder for you.

Next steps

Now that you understand all three fundamental patterns:

💡 Key takeaway: Different vectorization strategies suit different performance requirements. Manual vectorization gives maximum control, while Mojo’s vectorize function generates the chunk loop and its remainder for you. Choose based on your specific performance needs and development constraints.