Multi-Stage Pipeline Coordination

Overview

Implement a kernel that processes an image through a coordinated 3-stage pipeline where different thread groups handle specialized processing stages, synchronized with explicit barriers.

Note: You have specialized thread roles: Stage 1 (threads 0-127) loads and preprocesses data, Stage 2 (threads 128-255) applies blur operations, and Stage 3 (all threads) performs final smoothing.

Algorithm architecture: This puzzle implements a producer-consumer pipeline where different thread groups execute completely different algorithms within a single GPU block. Unlike traditional GPU programming where all threads execute the same algorithm on different data, this approach divides threads by functional specialization.

Pipeline concept: The algorithm processes data through three distinct stages, where each stage has specialized thread groups that execute different algorithms. Each stage produces data that the next stage consumes, creating explicit producer-consumer relationships that must be carefully synchronized with barriers.

Data dependencies and synchronization: Each stage produces data that the next stage consumes:

• Stage 1 → Stage 2: First stage produces preprocessed data for blur processing
• Stage 2 → Stage 3: Second stage produces blur results for final smoothing
• Barriers prevent race conditions by ensuring complete stage completion before dependent stages begin

Concretely, the multi-stage pipeline implements a coordinated image processing algorithm with three mathematical operations:

Stage 1 - Preprocessing Enhancement:

\[P[i] = I[i] \times 1.1\]

where \(P[i]\) is the preprocessed data and \(I[i]\) is the input data.

Stage 2 - Horizontal Blur Filter:

\[B[i] = \frac{1}{N_i} \sum_{k=-2}^{2} P[i+k] \quad \text{where } i+k \in [0, 255]\]

where \(B[i]\) is the blur result, and \(N_i\) is the count of valid neighbors within the tile boundary.

Stage 3 - Cascading Neighbor Smoothing:

\[F[i] = \begin{cases} (B[i] + B[i+1]) \times 0.6 & \text{if } i = 0 \\ ((B[i] + B[i-1]) \times 0.6 + B[i+1]) \times 0.6 & \text{if } 0 < i < 255 \\ (B[i] + B[i-1]) \times 0.6 & \text{if } i = 255 \end{cases}\]

where \(F[i]\) is the final output with cascading smoothing applied.

Thread Specialization:

• Threads 0-127: Compute \(P[i]\) for \(i \in \{0, 1, 2, \ldots, 255\}\) (2 elements per thread)
• Threads 128-255: Compute \(B[i]\) for \(i \in \{0, 1, 2, \ldots, 255\}\) (2 elements per thread)
• All 256 threads: Compute \(F[i]\) for \(i \in \{0, 1, 2, \ldots, 255\}\) (1 element per thread)

Synchronization Points:

\[\text{barrier}_1 \Rightarrow P[i] \text{ complete} \Rightarrow \text{barrier}_2 \Rightarrow B[i] \text{ complete} \Rightarrow \text{barrier}_3 \Rightarrow F[i] \text{ complete}\]

Key concepts

In this puzzle, you’ll learn about:

• Implementing thread role specialization within a single GPU block
• Coordinating producer-consumer relationships between processing stages
• Using barriers to synchronize between different algorithms (not just within the same algorithm)

The key insight is understanding how to design multi-stage pipelines where different thread groups execute completely different algorithms, coordinated through strategic barrier placement.

Why this matters: Most GPU tutorials teach barrier usage within a single algorithm - synchronizing threads during reductions or shared memory operations. But real-world GPU algorithms often require architectural complexity with multiple distinct processing stages that must be carefully orchestrated. This puzzle demonstrates how to transform monolithic algorithms into specialized, coordinated processing pipelines.

Previous vs. current barrier usage:

• Previous puzzles (P8, P12, P15): All threads execute the same algorithm, barriers sync within algorithm steps
• This puzzle: Different thread groups execute different algorithms, barriers coordinate between different algorithms

Thread specialization architecture: Unlike data parallelism where threads differ only in their data indices, this puzzle implements algorithmic parallelism where threads execute fundamentally different code paths based on their role in the pipeline.

Configuration

System parameters:

• Image size: SIZE = 1024 elements (1D for simplicity)
• Threads per block: TPB = 256 threads organized as (256, 1) block dimension
• Grid configuration: (4, 1) blocks to process entire image in tiles (4 blocks total)
• Data type: DType.float32 for all computations

Thread specialization architecture:

• Stage 1 threads: STAGE1_THREADS = 128 (threads 0-127, first half of block)

  • Responsibility: Load input data from global memory and apply preprocessing
  • Work distribution: Each thread processes 2 elements for efficient load balancing
  • Output: Populates input_shared[256] with preprocessed data

• Stage 2 threads: STAGE2_THREADS = 128 (threads 128-255, second half of block)

  • Responsibility: Apply horizontal blur filter on preprocessed data
  • Work distribution: Each thread processes 2 blur operations
  • Output: Populates blur_shared[256] with blur results

• Stage 3 threads: All 256 threads collaborate

  • Responsibility: Final smoothing and output to global memory
  • Work distribution: One-to-one mapping (thread i processes element i)
  • Output: Writes final results to global output array

Code to complete

alias TPB = 256  # Threads per block for pipeline stages
alias SIZE = 1024  # Image size (1D for simplicity)
alias BLOCKS_PER_GRID = (4, 1)
alias THREADS_PER_BLOCK = (TPB, 1)
alias dtype = DType.float32
alias layout = Layout.row_major(SIZE)

# Multi-stage processing configuration
alias STAGE1_THREADS = TPB // 2
alias STAGE2_THREADS = TPB // 2
alias BLUR_RADIUS = 2


fn multi_stage_image_blur_pipeline[
    layout: Layout
](
    output: LayoutTensor[mut=True, dtype, layout],
    input: LayoutTensor[mut=False, dtype, layout],
    size: Int,
):
    """Multi-stage image blur pipeline with barrier coordination.

    Stage 1 (threads 0-127): Load input data and apply 1.1x preprocessing
    Stage 2 (threads 128-255): Apply 5-point blur with BLUR_RADIUS=2
    Stage 3 (all threads): Final neighbor smoothing and output
    """

    # Shared memory buffers for pipeline stages
    input_shared = tb[dtype]().row_major[TPB]().shared().alloc()
    blur_shared = tb[dtype]().row_major[TPB]().shared().alloc()

    global_i = block_dim.x * block_idx.x + thread_idx.x
    local_i = thread_idx.x

    # Stage 1: Load and preprocess (threads 0-127)

    # FILL ME IN (roughly 10 lines)

    barrier()  # Wait for Stage 1 completion

    # Stage 2: Apply blur (threads 128-255)

    # FILL ME IN (roughly 25 lines)

    barrier()  # Wait for Stage 2 completion

    # Stage 3: Final smoothing (all threads)

    # FILL ME IN (roughly 7 lines)

    barrier()  # Ensure all writes complete

View full file: problems/p29/p29.mojo

Running the code

To test your solution, run the following command in your terminal:

uv run poe p29 --multi-stage

pixi run p29 --multi-stage

After completing the puzzle successfully, you should see output similar to:

Puzzle 29: GPU Synchronization Primitives
==================================================
TPB: 256
SIZE: 1024
STAGE1_THREADS: 128
STAGE2_THREADS: 128
BLUR_RADIUS: 2

Testing Puzzle 29A: Multi-Stage Pipeline Coordination
============================================================
Multi-stage pipeline blur completed
Input sample: 0.0 1.01 2.02
Output sample: 1.6665002 2.3331003 3.3996604
✅ Multi-stage pipeline coordination test PASSED!

Solution

The key insight is recognizing this as a pipeline architecture problem with thread role specialization:

1. Design stage-specific thread groups: Divide threads by function, not just by data
2. Implement producer-consumer chains: Stage 1 produces for Stage 2, Stage 2 produces for Stage 3
3. Use strategic barrier placement: Synchronize between different algorithms, not within the same algorithm
4. Optimize memory access patterns: Ensure coalesced reads and efficient shared memory usage