Puzzle 31: GPU Occupancy Optimization

Why this puzzle matters

Building on Puzzle 30: You’ve just learned GPU profiling tools and discovered how memory access patterns can create dramatic performance differences. Now you’re ready for the next level: resource optimization.

The Learning Journey:

• Puzzle 30 taught you to diagnose performance problems using Nsight profiling (nsys and ncu)
• Puzzle 31 teaches you to predict and control performance through resource management
• Together, they give you the complete toolkit for GPU optimization

What You’ll Discover: GPU performance isn’t just about algorithmic efficiency - it’s about how your code uses limited hardware resources. Every GPU has finite registers, shared memory, and execution units. Understanding occupancy - the ratio of active warps to maximum possible warps per SM - is crucial for:

• Latency hiding: Keeping the GPU busy while waiting for memory
• Resource allocation: Balancing registers, shared memory, and thread blocks
• Performance prediction: Understanding bottlenecks before they happen
• Optimization strategy: Knowing when to focus on occupancy vs other factors

Why This Matters Beyond GPUs: The principles you learn here apply to any parallel computing system where resources are shared among many execution units—from CPUs with hyperthreading to distributed computing clusters.

Overview

GPU Occupancy is the ratio of active warps to the maximum possible warps per SM. It determines how well your GPU can hide memory latency through warp switching.

SAXPY is a mnemonic for Single-precision Alpha times X plus Y. This puzzle explores three SAXPY kernels (y[i] = alpha * x[i] + y[i]) with nearly equivalent math but very different resource usage:

comptime SIZE = 32 * 1024 * 1024  # 32M elements - larger workload to show occupancy effects
comptime THREADS_PER_BLOCK = (1024, 1)
comptime BLOCKS_PER_GRID = (SIZE // 1024, 1)
comptime dtype = DType.float32
comptime layout = row_major[SIZE]()
comptime LayoutType = type_of(layout)
comptime ALPHA = Scalar[dtype](2.5)  # SAXPY coefficient


def minimal_kernel(
    y: TileTensor[mut=True, dtype, LayoutType, MutAnyOrigin],
    x: TileTensor[mut=False, dtype, LayoutType, ImmutAnyOrigin],
    alpha: Float32,
    size_dev: Int32,
):
    """Minimal SAXPY kernel - simple and register-light for high occupancy."""
    var size = Int(size_dev)
    var i = block_dim.x * block_idx.x + thread_idx.x
    if i < size:
        # Direct computation: y[i] = alpha * x[i] + y[i]
        # Uses minimal registers (~8), no shared memory
        y[i] = alpha * x[i] + y[i]

View full file: problems/p31/p31.mojo

def sophisticated_kernel(
    y: TileTensor[mut=True, dtype, LayoutType, MutAnyOrigin],
    x: TileTensor[mut=False, dtype, LayoutType, ImmutAnyOrigin],
    alpha: Float32,
    size_dev: Int32,
):
    """Sophisticated SAXPY kernel - over-engineered with excessive resource usage.
    """
    var size = Int(size_dev)
    # Maximum shared memory allocation (close to 48KB limit)
    var shared_cache = stack_allocation[
        dtype=dtype, address_space=AddressSpace.SHARED
    ](
        row_major[1024 * 12]()
    )  # 48KB

    var i = block_dim.x * block_idx.x + thread_idx.x
    var local_i = thread_idx.x

    if i < size:
        # REAL computational work that can't be optimized away - affects final result
        var base_x = x[i]
        var base_y = y[i]

        # Simulate "precision enhancement" - multiple small adjustments that add up
        # Each computation affects the final result so compiler can't eliminate them
        # But artificially increases register pressure
        var precision_x1 = base_x * 1.0001
        var precision_x2 = precision_x1 * 0.9999
        var precision_x3 = precision_x2 * 1.000001
        var precision_x4 = precision_x3 * 0.999999

        var precision_y1 = base_y * 1.000005
        var precision_y2 = precision_y1 * 0.999995
        var precision_y3 = precision_y2 * 1.0000001
        var precision_y4 = precision_y3 * 0.9999999

        # Multiple alpha computations for "stability" - should equal alpha
        var alpha1 = alpha * 1.00001 * 0.99999
        var alpha2 = alpha1 * 1.000001 * 0.999999
        var alpha3 = alpha2 * 1.0000001 * 0.9999999
        var alpha4 = alpha3 * 1.00000001 * 0.99999999

        # Complex polynomial "optimization" - creates register pressure
        var x_power2 = precision_x4 * precision_x4
        var x_power3 = x_power2 * precision_x4
        var x_power4 = x_power3 * precision_x4
        var x_power5 = x_power4 * precision_x4
        var x_power6 = x_power5 * precision_x4
        var x_power7 = x_power6 * precision_x4
        var x_power8 = x_power7 * precision_x4

        # "Advanced" mathematical series that contributes tiny amount to result
        var series_term1 = x_power2 * 0.0000001  # x^2/10M
        var series_term2 = x_power4 * 0.00000001  # x^4/100M
        var series_term3 = x_power6 * 0.000000001  # x^6/1B
        var series_term4 = x_power8 * 0.0000000001  # x^8/10B
        var series_correction = (
            series_term1 - series_term2 + series_term3 - series_term4
        )

        # Over-engineered shared memory usage with multiple caching strategies
        if local_i < 1024:
            shared_cache[local_i] = precision_x4
            shared_cache[local_i + 1024] = precision_y4
            shared_cache[local_i + 2048] = alpha4
            shared_cache[local_i + 3072] = series_correction
        barrier()

        # Load from shared memory for "optimization"
        var cached_x = shared_cache[local_i] if local_i < 1024 else precision_x4
        var cached_y = (
            shared_cache[local_i + 1024] if local_i < 1024 else precision_y4
        )
        var cached_alpha = (
            shared_cache[local_i + 2048] if local_i < 1024 else alpha4
        )
        var cached_correction = (
            shared_cache[local_i + 3072] if local_i
            < 1024 else series_correction
        )

        # Final "high precision" computation - all work contributes to result
        var high_precision_result = (
            cached_alpha * cached_x + cached_y + cached_correction
        )

        # Over-engineered result with massive resource usage but mathematically ~= alpha*x + y
        y[i] = high_precision_result

View full file: problems/p31/p31.mojo

def balanced_kernel(
    y: TileTensor[mut=True, dtype, LayoutType, MutAnyOrigin],
    x: TileTensor[mut=False, dtype, LayoutType, ImmutAnyOrigin],
    alpha: Float32,
    size_dev: Int32,
):
    """Balanced SAXPY kernel - efficient optimization with moderate resources.
    """
    var size = Int(size_dev)
    # Reasonable shared memory usage for effective caching (16KB)
    var shared_cache = stack_allocation[
        dtype=dtype, address_space=AddressSpace.SHARED
    ](
        row_major[1024 * 4]()
    )  # 16KB total

    var i = block_dim.x * block_idx.x + thread_idx.x
    var local_i = thread_idx.x

    if i < size:
        # Moderate computational work that contributes to result
        var base_x = x[i]
        var base_y = y[i]

        # Light precision enhancement - less than sophisticated kernel
        var enhanced_x = base_x * 1.00001 * 0.99999
        var enhanced_y = base_y * 1.00001 * 0.99999
        var stable_alpha = alpha * 1.000001 * 0.999999

        # Moderate computational optimization
        var x_squared = enhanced_x * enhanced_x
        var optimization_hint = x_squared * 0.000001

        # Efficient shared memory caching - only what we actually need
        if local_i < 1024:
            shared_cache[local_i] = enhanced_x
            shared_cache[local_i + 1024] = enhanced_y
        barrier()

        # Use cached values efficiently
        var cached_x = shared_cache[local_i] if local_i < 1024 else enhanced_x
        var cached_y = (
            shared_cache[local_i + 1024] if local_i < 1024 else enhanced_y
        )

        # Balanced computation - moderate work, good efficiency
        var result = stable_alpha * cached_x + cached_y + optimization_hint

        # Balanced result with moderate resource usage (~15 registers, 16KB shared)
        y[i] = result

View full file: problems/p31/p31.mojo

Your task

Use profiling tools to investigate three kernels and answer analysis questions about occupancy optimization. The kernels compute equivalent results (to within test tolerance) but use resources very differently - your job is to discover why performance and occupancy behave counterintuitively!

The worked occupancy calculation in this puzzle is based on NVIDIA A10G (Ampere 8.6) hardware; the measured resource and timing figures come from a B200 (Blackwell 10.0). Your results will vary depending on your GPU vendor and architecture (NVIDIA: Pascal/Turing/Ampere/Ada/Hopper/Blackwell, AMD: RDNA/GCN, Apple: M1/M2/M3/M4/M5), but the fundamental concepts, methodology, and insights remain universally applicable across modern GPUs. Use pixi run gpu-specs to get your specific hardware values.

Hardware is not the only variable — the build is too. Register counts and timings differ between a --debug-level=full build and a release one, and occupancy follows the registers. Shared memory per block does not move. Any number below that you intend to compare against your own must say which build it came from; see the debug-build trade-off.

Configuration

Requirements:

• NVIDIA GPU with CUDA toolkit
• Nsight Compute from Puzzle 30

⚠️ GPU compatibility note: The default configuration uses aggressive settings that may fail on older or lower-capability GPUs:

comptime SIZE = 32 * 1024 * 1024  # 32M elements (~128MB per array)
comptime THREADS_PER_BLOCK = (1024, 1)  # 1024 threads per block
comptime BLOCKS_PER_GRID = (SIZE // 1024, 1)  # 32768 blocks

If you encounter launch failures, reduce these values in problems/p31/p31.mojo:

• For older GPUs: Use THREADS_PER_BLOCK = (512, 1) and SIZE = 16 * 1024 * 1024
• For limited memory GPUs (< 2GB): Use SIZE = 8 * 1024 * 1024 or SIZE = 4 * 1024 * 1024
• For grid dimension limits: The BLOCKS_PER_GRID will automatically adjust with SIZE

Occupancy Formula:

Theoretical Occupancy = min(
    Registers Per SM / (Registers Per Thread × Threads Per Block),
    Shared Memory Per SM / Shared Memory Per Block,
    Max Blocks Per SM
) × Threads Per Block / Max Threads Per SM

The investigation

Step 1: Test the kernels

pixi shell -e nvidia
mojo problems/p31/p31.mojo --all

All three compute the same SAXPY result to within test tolerance—the sophisticated and balanced kernels add small correction terms, so the tests compare with rtol=1e-3 and rtol=1e-4 rather than exactly. The mystery: why do they have different performance?

Step 2: Benchmark performance

mojo problems/p31/p31.mojo --benchmark

Record the reported time for each kernel. Step 4 is where you connect those numbers to the resources each kernel consumes.

Step 3: Build for profiling

mojo build problems/p31/p31.mojo -o problems/p31/p31_profiler

No --debug-level=full here, unlike the profiling walkthrough in Puzzle 30. This puzzle’s subject is resource usage, and debug metadata inflates register counts — enough, on a B200, to drag the sophisticated kernel’s measured occupancy from 82% down to 50% and invert the comparison you are about to make. Profile a release build whenever the numbers you want are the resource numbers.

Step 4: Profile resource usage

# Profile each kernel's resource usage
ncu --set basic --section=LaunchStats problems/p31/p31_profiler --minimal
ncu --set basic --section=LaunchStats problems/p31/p31_profiler --sophisticated
ncu --set basic --section=LaunchStats problems/p31/p31_profiler --balanced

Record the resource usage for occupancy analysis.

Step 5: Calculate theoretical occupancy

First, identify your GPU architecture and detailed specs:

pixi run gpu-specs

Note: gpu-specs automatically detects your GPU vendor (NVIDIA/AMD/Apple) and shows all architectural details derived from your hardware - no lookup tables needed!

Common Architecture Specs (Reference):

Note that the two Ampere compute capabilities differ: the data-center 8.0 part (A100) has 164KB of shared memory and 2,048 threads per SM, while the 8.6 parts (RTX 30xx, A10G) have 100KB and 1,536. That distinction drives the worked solution below.

📚 Official Documentation:

• NVIDIA CUDA Compute Capability Table
• CUDA Programming Guide - Compute Capabilities
• Hopper Architecture In-Depth
• Ampere Architecture Whitepaper

⚠️ Note: These are theoretical maximums. Actual occupancy may be lower due to hardware scheduling constraints, driver overhead, and other factors.

Using your GPU specs and the occupancy formula:

• Threads Per Block: 1024 (from our kernel)

Use the occupancy formula and your hardware specifications to predict each kernel’s theoretical occupancy.

Step 6: Measure actual occupancy

# Measure actual occupancy for each kernel
ncu --metrics=smsp__warps_active.avg.pct_of_peak_sustained_active problems/p31/p31_profiler --minimal
ncu --metrics=smsp__warps_active.avg.pct_of_peak_sustained_active problems/p31/p31_profiler --sophisticated
ncu --metrics=smsp__warps_active.avg.pct_of_peak_sustained_active problems/p31/p31_profiler --balanced

Compare the actual measured occupancy with your theoretical calculations - this is where the mystery reveals itself!

Key insights

💡 Occupancy Threshold: Once you have sufficient occupancy for latency hiding (~25-50%), additional occupancy provides diminishing returns.

💡 Memory Bound vs Compute Bound: SAXPY is memory-bound. Memory bandwidth often matters more than occupancy for memory-bound kernels.

💡 Resource Efficiency: Modern GPUs can handle moderate register pressure (20-40 registers/thread) without dramatic occupancy loss.

Your task: Answer the following questions

After completing the investigation steps above, answer these analysis questions to solve the occupancy mystery:

Performance Analysis (Step 2):

1. Which kernel is fastest? Which is slowest? Record the timing differences.

Resource Profiling (Step 4):

2. Record for each kernel: Registers Per Thread, Shared Memory Per Block, Warps Per SM

Theoretical Calculations (Step 5):

3. Calculate theoretical occupancy for each kernel using your GPU specs and the occupancy formula. Which should be highest/lowest?

Measured Occupancy (Step 6):

4. How do the measured occupancy values compare to your calculations?

The Occupancy Mystery:

5. Why do all three kernels land in a similar occupancy band (the exact figures vary by GPU architecture) despite dramatically different resource usage?
6. Why does performance vary so little — about 14% between fastest and slowest on a B200 — when shared memory usage varies from 0KB to 49KB?
7. What does this reveal about the relationship between theoretical occupancy calculations and real-world GPU behavior?
8. For this SAXPY workload, what is the actual performance bottleneck if it’s not occupancy?

Solution

Register Limit = 65,536 / (19 × 1,024) = 3.36 blocks per SM
Shared Memory Limit = 100KB / 0KB = ∞ blocks per SM
Hardware Block Limit = 16 blocks per SM

Thread Limit = 1,536 / 1,024 = 1 block per SM (floor)
Actual Blocks = min(3, ∞, 1) = 1 block per SM
Theoretical Occupancy = (1 × 1,024) / 1,536 = 66.7%

Register Limit = 65,536 / (25 × 1,024) = 2.56 blocks per SM
Shared Memory Limit = 100KB / 16.4KB = 6.10 blocks per SM
Hardware Block Limit = 16 blocks per SM

Thread Limit = 1,536 / 1,024 = 1 block per SM (floor)
Actual Blocks = min(2, 6, 1) = 1 block per SM
Theoretical Occupancy = (1 × 1,024) / 1,536 = 66.7%

Register Limit = 65,536 / (40 × 1,024) = 1.6 blocks per SM
Shared Memory Limit = 100KB / 49.2KB = 2.03 blocks per SM
Hardware Block Limit = 16 blocks per SM

Thread Limit = 1,536 / 1,024 = 1 block per SM (floor)
Actual Blocks = min(1, 2, 1) = 1 block per SM
Theoretical Occupancy = (1 × 1,024) / 1,536 = 66.7%

Single Kernel Pass Analysis:
- Input arrays: 32M × 4 bytes × 2 arrays = 256MB read
- Output array: 32M × 4 bytes × 1 array = 128MB write
- Total per kernel: 384MB memory traffic

Measured (B200, minimal kernel): 0.1399ms per iteration
Implied bandwidth: 384MB / 0.1399ms ≈ 2.7 TB/s

# Find your GPU specs
pixi run gpu-specs

# Measure resource consumption
ncu --set basic --section=LaunchStats your_kernel

# Measure achieved occupancy
ncu --metrics=smsp__warps_active.avg.pct_of_peak_sustained_active your_kernel

# Always validate with actual performance measurements
ncu --set roofline --section=MemoryWorkloadAnalysis your_kernel

OCCUPANCY ANALYSIS → OPTIMIZATION STRATEGY:

High occupancy (>70%) + Good performance:
→ Occupancy is sufficient, focus on other bottlenecks

Low occupancy (<30%) + Poor performance:
→ Increase occupancy through resource optimization

Good occupancy (50-70%) + Poor performance:
→ Look for memory bandwidth, cache, or computational bottlenecks

Low occupancy (<30%) + Good performance:
→ Workload doesn't need high occupancy (memory-bound)