Quick Start

Activate Environment

conda activate gpu-jit-opt

Run Phase 0 (Baseline Measurement)

Measure kernel performance across different compiler knobs:

python experiments/phase0_baseline_table.py

Output: results/tables/phase0_baseline.csv

Time: ~3 seconds

Run Phase 3 Rollout Logging (NVML-only)

Generate environment step logs with lightweight telemetry:

python phase3_rollout_log.py --use-nvml --kernels gemm reduction softmax --matrix-sizes 256 512 --episodes-per-case 1 --max-steps 10 --warmup 1 --repeats 5

Output: results/tables/phase3_rollout.csv, phase3_episode_summary.csv

Time: ~2-5 minutes

Train PPO Agent

Train a reinforcement learning agent to learn optimal configurations:

python train_rl.py --total-steps 50000 --max-episode-len 50 --use-nvml --eval-freq 5000 --n-eval-episodes 5

Output: results/models/<run_tag>.zip (exact path is recorded in the training summary JSON)

Time: ~15-30 minutes (NVML-only)

Next Steps

1. Review Phase 0 baseline results in CSV
2. Validate kernel correctness with Phase 2 tests
3. Run Phase 3 environment logging
4. Train or load a PPO policy
5. Analyze results and visualize improvements

Installation & Setup

Prerequisites

• NVIDIA GPU with CUDA compute capability 6.0+
• NVIDIA CUDA Toolkit 11.8+ installed
• Nsight Compute (for Phase 1 counter collection)
• Conda or Miniconda

Step 1: Clone and Setup Environment

# Create conda environment
conda create -n gpu-jit-opt python=3.10

# Activate
conda activate gpu-jit-opt

# Install PyTorch with CUDA support
conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia

# Install Numba CUDA
conda install numba

# Install remaining dependencies
pip install -r requirements.txt

Step 2: Verify Installation

python check_cuda.py

Should report:

• CUDA available: True
• Your GPU name and memory
• CUDA operations work

Step 3: Verify Kernels

pytest tests/test_kernels.py -v

Should show all tests passing.

Troubleshooting Installation

Core Concepts: Registers & Occupancy

GPU Memory Hierarchy

GPUs organize memory in multiple tiers with different speeds and scopes:

• Registers: Fastest, per-thread storage inside GPU core. Limited count per thread.
• Shared Memory: Fast memory shared by threads in a block. Enables cooperation.
• L2 Cache: On-chip cache for all accesses. Improves locality.
• Global Memory (VRAM): Largest but slowest. GPU main memory.

Register Pressure

Register count per thread is a critical resource constraint:

• Low register usage: More threads can be resident → higher occupancy. But may require "spilling" to slower memory.
• High register usage: Few threads resident → lower occupancy. But fewer memory transactions.
• Trade-off: The optimal configuration depends on kernel compute intensity and memory access patterns.

What is "Occupancy"?

Occupancy is the ratio of active warps to maximum possible warps on a Streaming Multiprocessor (SM):

occupancy = active_warps_per_SM / max_warps_per_SM

• On Ampere (SM 8.6), max warps per SM is 48
• So occupancy of 0.75 means 36 warps are resident
• Occupancy of 1.0 means 48 warps (100% occupancy)

GPU Architecture Definitions

• Thread: Smallest unit of execution.
• Warp: Group of 32 threads that execute in lockstep (SIMT model). NVIDIA GPUs process 32 threads together.
• Block: User-defined group of threads that can cooperate via shared memory and synchronization.
• Grid: All blocks launched for a single kernel call.
• SM (Streaming Multiprocessor): Physical "core cluster" on the GPU. Each SM runs blocks and manages warps.

Why Occupancy Matters

• Memory-bound kernels: Higher occupancy helps hide memory latency. While some warps wait for data, others execute → better throughput.
• Compute-bound kernels: Occupancy is less critical. The bottleneck is compute units, not memory latency.
• Important caveat: 100% occupancy does NOT guarantee fast execution. A kernel may have 100% occupancy but still be slow if it's memory-bound and has poor data reuse.

How Occupancy Changes

Occupancy is limited by GPU resource constraints:

• Registers per thread: More registers → fewer warps fit → lower occupancy
• Threads per block: Larger blocks → fewer blocks fit per SM → may reduce occupancy
• Shared memory per block: More shared memory per block → fewer blocks fit

Register Capping Trade-offs

We use compiler flags to limit register usage:

reg_cap = 32  # Request: compile to use ≤ 32 registers per thread
# Applied via Numba: @cuda.jit(..., max_registers=32)

If you cap registers too low, the compiler must spill excess values into local memory (which lives in global memory), making the kernel much slower. But allowing many registers means fewer threads can be resident on an SM (lower occupancy). The optimal register cap depends on the kernel's compute intensity and memory access patterns.

The Register Cap vs Occupancy Trade-off (Example)

For a kernel with 256 threads/block on RTX 3050 Ti:

Compiler Knobs & Configuration

Key Knobs in This Project

1. Block Size (Threads per Block)

Determines how many threads cooperate in a block.

• Values: 64, 128, 256, 512
• Effect on occupancy: Larger blocks → fewer blocks fit per SM → may reduce occupancy.
• Effect on performance: Larger blocks can improve shared memory efficiency and reduce launch overhead.

2. Register Cap (max_registers)

Maximum registers per thread requested during compilation.

• Values: 0 (default), 32, 64
• Effect: Lower caps force the compiler to spill, reducing register usage but potentially increasing memory pressure.
• Numba implementation: @cuda.jit(..., max_registers=32)

Why We Sweep Over These Knobs

Different kernels respond differently:

• Compute-heavy kernels (GEMM) may benefit from high register budgets and larger blocks.
• Memory-heavy kernels (reduction) may be sensitive to occupancy and prefer smaller register caps.
• Mixed kernels (softmax) show complex trade-offs.

GPU Metrics & Profiling

CUPTI Metrics (Nsight Compute)

Hardware counters collected via NVIDIA Nsight Compute CLI (ncu). These are real measurements taken while the kernel ran:

Achieved Occupancy

What it measures: Proxy for warps active during execution.

Metric: sm__warps_active.avg.pct_of_peak_sustained_active

• Range: 0–1 (or 0–100%)
• Interpretation: Fraction of max warps that were on average active during the kernel run.
• Example: 0.6 on RTX 3050 Ti means ~29 of 48 max warps per SM were active on average.
• Why it matters: High occupancy helps latency hiding for memory-bound kernels.

L2 Cache Hit Rate

What it measures: Fraction of memory requests served from L2 cache (vs DRAM).

Metric: lts__t_sector_hit_rate.pct

• Range: 0–1
• Interpretation: Higher = better locality, fewer expensive DRAM transactions.
• Example: 0.9 hit rate means 90% of memory accesses were served from L2, only 10% needed DRAM.
• Why it matters: Low hit rates indicate streaming memory access (poor temporal locality).

DRAM Bandwidth

What it measures: DRAM throughput as percentage of peak bandwidth.

• Range: 0–1
• Interpretation: High values (0.7–1.0) suggest memory-bound kernel pushing max bandwidth.
• Example on RTX 3050 Ti: Peak DRAM BW ~ 288 GB/s. If util is 0.5, actual BW ~ 144 GB/s.
• Why it matters: If DRAM BW is high and speedup is low, kernel is memory-limited.

SM Active Percentage

Compute unit utilization. Metric: sm__throughput.avg.pct_of_peak_sustained_active

• Range: 0–1
• Interpretation: Fraction of time SMs were doing useful work (not stalled).

NVML Telemetry (Lightweight, Always Available)

Collected via pynvml, no special permissions usually needed:

GPU Utilization

How busy the GPU compute is. Windowed metric on Windows.

Memory Utilization

How busy the memory interface is.

Memory Used Fraction

VRAM currently in use / total VRAM.

Temperature

GPU core temperature in Celsius.

Phase 0: Foundational Baseline Table

Purpose

Measure raw kernel performance across a systematic sweep of compiler knobs without profiling to establish baseline metrics. Phase 0 is fast and doesn't require special permissions or profiling tools.

What Phase 0 Measures

A systematic sweep over:

• Kernel types: GEMM (compute-heavy), Reduction (memory-bound), Softmax (mixed)
• Problem sizes: 256, 512, 1024
• Block sizes: 64, 128, 256 threads per block
• Register caps: default (0), 32, 64
• Total configurations: 3 × 3 × 3 × 3 = 81 unique runs

Each configuration is timed multiple times, and mean/std/min/max are recorded.

Why Phase 0 Matters

Phase 0 answers three critical questions:

1. "Do compiler knobs actually matter?" By varying register caps and block sizes, we establish whether runtime changes meaningfully.
2. "What are the trends?" Is GEMM faster with high register budgets? Does Reduction prefer lower registers?
3. "What's a good baseline?" The Phase 0 results provide a comparison point for RL (Phase 3) improvements.

Output: phase0_baseline.csv

How to Run Phase 0

python experiments/phase0_baseline_table.py

Expected Output

• Total configurations: 81
• Estimated time: ~3 seconds
• Pretty table showing best configs per kernel
• CSV saved to results/tables/phase0_baseline.csv

Interpretation Guide

Key Insight: Theoretical occupancy (theor_occ) is calculated from registers and block size, but actual runtime depends on compute intensity, memory patterns, and other factors.

Per-Kernel Interpretation

• GEMM (compute-heavy): Best time often occurs with higher register usage (less spilling), even if occupancy is lower. Signal: high SM active %, low memory bandwidth.
• Reduction (memory-bound): Higher occupancy often correlates with lower time. Memory bandwidth is crucial. Signal: depends on access patterns and cache hierarchy.
• Softmax (mixed): Complex trade-offs; empirical results vary by matrix size and block size.

Why Some Rows Have Huge Timing Noise

A few common reasons:

• CPU-side timing jitter on Windows: WDDM driver model can introduce delays.
• Small number of repeats: With only 3 repeats, one slow outlier inflates mean and std.
• Under-utilization: Very small grids underutilize GPU, so overheads dominate.

Sanity check: Prefer interpreting rows with small std. If noise is high, increase repeats or run on a quieter system.

Phase 1: Hardware Counters via Nsight Compute

Purpose

Collect detailed hardware performance metrics to understand GPU behavior during kernel execution.

Prerequisites

• NVIDIA Nsight Compute installed (typically with CUDA Toolkit)
• Administrator privileges on Windows (for GPU counter access)
• Phase 0 and Phase 2 tests passing

How to Run Phase 1

On Windows, run Command Prompt as Administrator:

python experiments/phase1_collect_counters.py

Output: phase1_result.csv

Similar structure to Phase 0, plus hardware metrics:

Permutation Errors

Profiling Overhead

Phase 1 is significantly slower than Phase 0 due to profiler overhead. A single configuration may take 5–30 seconds.

Phase 2: Kernel Correctness Validation

Purpose

Ensure all kernels compute correct results before using them for performance measurement or RL training.

Two-Tier Testing Strategy

Tier 1: Small Correctness Suite (default)

Fast tests on small/edge-case sizes:

pytest tests/test_kernels.py -v

Time: ~5–10 seconds

Coverage: Boundary conditions, small matrices, various block sizes

Tier 2: Performance-Scale Suite (opt-in)

Tests on benchmark-like sizes:

pytest tests/test_kernels.py -v --runslow

Time: ~30–60 seconds

Coverage: Sizes 256–1024, realistic block-register combinations

Test Coverage

• GEMM Correctness: Output C matches A × B (within FP32 tolerance)
• Reduction Correctness: Output matches sum(x)
• Softmax Correctness: Each row sums to ~1.0, all values ≥ 0

Floating-Point Tolerance

GPU results may differ slightly from CPU due to:

• Different instruction order → accumulated rounding error
• Math library differences (e.g., exp, log)

Tests use rtol=1e-4, atol=1e-5 (relative and absolute tolerances).

Common Issues

Phase 3: RL Environment & PPO Training

Overview

Phase 3 wraps the kernels and metrics in a Gymnasium environment for reinforcement learning, enabling automated discovery of optimal configurations.

Beginner Concept: What is Reinforcement Learning?

Reinforcement Learning (RL) is a paradigm where an agent learns by trial and error:

• Agent observes the current state (e.g., kernel name, previous performance metrics)
• Agent takes an action (e.g., "try block_size=256, reg_cap=32")
• Environment responds with a new state and a reward (e.g., speedup score)
• Agent learns which actions lead to high rewards over time

In this project, the "best action" is the configuration that achieves maximum speedup over the baseline.

What is PPO (Proximal Policy Optimization)?

PPO is a popular RL algorithm that trains a policy network (a neural network) to predict good actions. Key advantages:

• Sample-efficient: Learns quickly even with limited data
• Stable: Uses clipping to prevent catastrophic updates
• Well-supported: Stable-Baselines3 provides industrial-strength implementation

We do NOT implement PPO from scratch; we use Stable-Baselines3's vetted implementation.

Environment Structure

Action Space

MultiDiscrete actions over (block_size, reg_cap):

action_space = MultiDiscrete([num_block_sizes, num_reg_caps]) = MultiDiscrete([3, 3])
# block_size_idx ∈ {0,1,2} → [64, 128, 256]
# reg_cap_idx   ∈ {0,1,2} → [0, 32, 64]

Observation Space

Normalized bounded vector [0, 1]. The agent observes:

- Kernel one-hot: [1, 0, 0] for gemm, [0, 1, 0] for reduction, [0, 0, 1] for softmax
- Previous action: [block_size_idx/2, reg_cap_idx/2] (normalized)
- CUPTI metrics: 4 values [0,1] (occupancy, L2, DRAM BW, SM active)
    (these 4 slots are always present; they are zeros when CUPTI collection is disabled/unavailable)
- NVML metrics: 4 values [0,1] (GPU util, mem util, mem used, temp)
Total dimension: 13 (with NVML enabled, default) or 9 (if NVML is disabled)

Why These Observations?

• Kernel one-hot: Different kernels respond differently to configs → agent learns kernel-specific policies
• Previous action: Feedback loop: agent sees its own previous choice → learns sequences
• CUPTI metrics: Hardware-level signals (rich but slow to collect)
• NVML metrics: Lightweight telemetry (always available, fast)

Normalization for ML

All observations are normalized to [0,1] to keep the observation space bounded and ML-friendly:

• Percentages (0-100%): Divided by 100 → [0,1]
• Ratios: Already [0,1]
• Temperature: Normalized by 100 → temp/100 clamped to [0,1]
• Out-of-range: Clamped to [0,1] for safety

Reward

baseline_ms = runtime with (block_size=256, reg_cap=0)
measured_ms = runtime with agent-chosen (block_size, reg_cap)
speedup = baseline_ms / measured_ms
reward = speedup - 1
# reward > 0: faster than baseline → positive feedback
# reward = 0: same speed as baseline
# reward < 0: slower than baseline → negative feedback

Reward Examples

• If baseline is 10ms and measured is 8ms: speedup = 10/8 = 1.25 → reward = 0.25 ✓ Good!
• If baseline is 10ms and measured is 12ms: speedup = 10/12 = 0.83 → reward = -0.17 ✗ Bad!
• If baseline is 10ms and measured is 10ms: speedup = 1.0 → reward = 0.0 (neutral)

Why Speedup as Reward?

This reward design has nice properties:

• Scale-invariant: A 10% speedup on fast kernels (0.1ms → 0.09ms) gets same reward as on slow ones (100ms → 90ms).
\n
• Intuitive: Human experts naturally think in terms of speedup, not raw milliseconds.
\n
• Bounded: Speedup is typically in realistic range (0.5x to 2x) so rewards don't explode.
\n

Episode Structure

1. Reset: Pick random kernel + matrix size, measure baseline with (block_size=256, reg_cap=0)
2. Steps: Agent chooses actions for max_steps iterations
3. Termination: Episode ends after max_steps

Running Phase 3: Rollout Logging

NVML-only (Fast, No Privileges)

python phase3_rollout_log.py --use-nvml --kernels gemm reduction softmax --matrix-sizes 256 512 --episodes-per-case 1 --max-steps 10 --warmup 1 --repeats 5

Time: ~2–5 minutes

CUPTI+NVML (Slow, Administrator Required)

python phase3_rollout_log.py --use-cupti --use-nvml --kernels gemm reduction softmax --matrix-sizes 256 512 --episodes-per-case 1 --max-steps 10 --cupti-timeout-s 180

Time: ~20–60 minutes

Running Phase 3: PPO Training

⚠️ CRITICAL: Performance Warning

NVML-only Training (RECOMMENDED)

python train_rl.py --total-steps 50000 --max-episode-len 50 --use-nvml --eval-freq 5000 --n-eval-episodes 5

Time: ~9–10 minutes on RTX 3050 Ti

Actual Result: 50,176 steps, 9.7 minutes, mean_reward=3.15, FPS=85

Why: Fast, stable, no admin needed. Achieves excellent convergence.

CUPTI Analysis (Optional, After Training)

Collect detailed metrics on short subset with your trained policy:

python phase3_rollout_log.py --use-cupti --use-nvml --kernels gemm reduction softmax --matrix-sizes 256 512 --episodes-per-case 3 --max-steps 10

Time: ~30–60 minutes

CUPTI+NVML Limited Training (Research-Grade, 3,000–5,000 steps)

Strategy for Publication-Ready Results:

# Option A: Quick validation (3,000 steps, 4–5 hours)
python train_rl.py --total-steps 3000 --max-episode-len 25 --batch-size 512 --use-cupti --use-nvml --cupti-timeout-s 180 --eval-freq 1500 --n-eval-episodes 2

# Option B: Full convergence (5,000 steps, 6–8 hours)
python train_rl.py --total-steps 5000 --max-episode-len 25 --batch-size 512 --use-cupti --use-nvml --cupti-timeout-s 180 --eval-freq 2500 --n-eval-episodes 3

Expected Times: 3,000 steps ≈ 4–5 hours | 5,000 steps ≈ 6–8 hours

What You Get: 13D observation space (vs 9D NVML-only) with rich hardware counters from CUPTI

Avoid >5,000 steps: Risk of CUDA context degradation. Post-hoc CUPTI analysis is better alternative.

Post-Training CUPTI Analysis (RECOMMENDED for Research):

# After NVML training completes, collect rich CUPTI metrics on small subset
python phase3_rollout_log.py --use-cupti --use-nvml --kernels gemm reduction softmax --matrix-sizes 256 512 1024 --episodes-per-case 3 --max-steps 10 --cupti-timeout-s 180

Time: ~30–60 minutes | Benefit: Publication-quality CUPTI data without training overhead

Key PPO Hyperparameters

Training Artifacts

• results/models/<run_tag>.zip — Final trained policy
• results/models/<run_tag>_*_steps.zip — Periodic checkpoints
• results/models/<run_tag>_best/best_model.zip — Best policy (if eval enabled)
• results/logs/<run_tag>/train_rl_<gpu_tag>.log — Training log
• results/logs/<run_tag>/training_summary_<gpu_tag>.json — Hyperparameters and artifact paths
• results/logs/tensorboard/ — TensorBoard events (run: tensorboard --logdir=results/logs/tensorboard)

NVML-only vs CUPTI+NVML Training: Which Should You Use?

Recommendation

\n
• Normal projects: Always use --use-nvml. Train in 15-30 min.
\n
• Publication/research: Train with NVML, then collect CUPTI metrics separately with phase3_rollout_log.py on smaller subset.
\n
• Do NOT: Combine CUPTI+full training (>5k steps). Will take days or crash.
\n

Windows WDDM Driver Limitations

Why is CUPTI+full training so slow/unstable on Windows?

\n
• WDDM driver model: Adds overhead to profiler calls
\n
• ncu profiling: Each kernel run must be profiled in isolation (can't overlap)
\n
• Context switching: Sustained profiling destabilizes CUDA contexts after hours
\n
• Memory fragmentation: 22+ hours of continuous profiling leaves GPU memory fragmented
\n

Workaround: Use NVML metrics during training (always available), profile with CUPTI afterwards on small subset.

Output Artifacts & CSV Schemas

Phase 0: phase0_baseline.csv

Row per configuration. Columns: kernel, matrix_size, block_size, threads_per_block, reg_cap, est_regs, actual_regs, theor_occ, time_ms_mean, time_ms_std, time_ms_min, time_ms_max, [achieved_occ (optional)]

Phase 1: phase1_result.csv

Similar to Phase 0, plus metric columns:

• ok (bool): success
• reason (string): error reason
• {metric}_raw, {metric}_norm for achieved_occupancy, l2_hit_rate, dram_bw_pct, sm_active_pct

Phase 3: phase3_rollout.csv (step-level)

Row per environment step:

• episode_id, episode_seed, step
• kernel, matrix_size, block_size, reg_cap
• time_ms, baseline_ms, speedup, reward
• cupti_ok, cupti_reason
• cupti_{metric}_norm (if enabled)
• nvml_{metric}_norm (if enabled)

Phase 3: phase3_episode_summary.csv (episode-level)

Row per episode:

• episode_id, episode_seed, kernel, matrix_size
• max_steps, use_cupti, use_nvml
• baseline_ms, mean_time_ms, mean_reward
• best_speedup, best_reward, steps

Phase 4a: Rollout Evaluation Results (Actual Data)

Configuration: 45 episodes × 20 steps = 900 step evaluations across 9 kernels (3 kernels × 3 matrix sizes) × 5 episodes per case

Agent: NVML-trained PPO (3.15x mean reward from 50,176 steps)

Key Finding: Softmax scales best with problem size (1.64x → 2.18x mean from 256→1024). GEMM modest (1.22x → 1.15x). Reduction highly variable but achievable (0.68–2.89x range).

Interpreting Results

Phase 0 Analysis

Basic sanity checks:

• All rows should have positive time_ms values
• Std should be small relative to mean (low timing jitter)
• Theoretical occupancy should match pattern with register usage

Python snippet:

import pandas as pd
df = pd.read_csv('results/tables/phase0_baseline.csv')
best = df.loc[df.groupby('kernel')['time_ms_mean'].idxmin()]
print(best[['kernel', 'matrix_size', 'block_size', 'reg_cap', 'time_ms_mean']])

Phase 1 Analysis

Correlate metrics with performance:

• High DRAM BW + Low L2 hit rate → memory-bound config
• High occupancy + Low SM active → might be memory stalled
• Low occupancy + High SM active → likely register-constrained

Phase 3 Results Interpretation

Best config per kernel:

import pandas as pd
df = pd.read_csv('results/tables/phase3_rollout.csv')
best = df.loc[df.groupby('kernel')['speedup'].idxmax()]
print(best[['kernel', 'matrix_size', 'block_size', 'reg_cap', 'speedup']])

Learning curves:

import matplotlib.pyplot as plt
episode_df = pd.read_csv('results/tables/phase3_episode_summary.csv')
plt.plot(episode_df['episode_id'], episode_df['best_speedup'], 'o-')
plt.xlabel('Episode')
plt.ylabel('Best Speedup')
plt.title('PPO Training Progress')
plt.show()

Common Gotchas

Kernel Implementations

GEMM (General Matrix Multiplication)

File: kernels/gemm.py

Compute Pattern: Tiled matrix multiplication using shared memory blocking.

• Workload: Multiplies two NxN matrices → compute-heavy
• Compute Intensity: High (O(N³) operations, O(N²) memory)
• Typical Behavior: Benefits from large register budgets and higher occupancy
• Block sizes: 64, 128, 256 (TILE parameter)

Reduction

File: kernels/reduction.py

Compute Pattern: Parallel sum reduction with block-level synchronization.

• Workload: Sums N² elements into a single output – memory-bound
• Compute Intensity: Low (O(N²) operations, O(N²) memory)
• Typical Behavior: Sensitive to occupancy; lower register usage helps
• Block sizes: 64, 128, 256

Softmax

File: kernels/softmax.py

Compute Pattern: Row-wise softmax: exp, normalization, log. Reduces NxN matrix row-by-row.

• Workload: Mixed compute + memory
• Compute Intensity: Medium (transcendental math, but memory-light per row)
• Typical Behavior: Complex trade-offs; varies by size and config
• Block sizes: 64, 128, 256

Profiling & Metrics Collection

CUDA Event Timing

Kernel execution timing via NVIDIA CUDA events:

start = cuda.event()
stop = cuda.event()
start.record()
kernel[grid, block](*args)
stop.record()
stop.synchronize()
elapsed_ms = start.time_till(stop)

Accuracy: Sub-microsecond on recent GPUs. Repeating and averaging stabilizes results.

CUPTI Collection

Nsight Compute CLI invoked via subprocess:

ncu --set full [--metrics specific_metrics] python run_kernel.py

Overhead: Significant (5–30 seconds per kernel).

Timeout handling: If ncu hangs, subprocess terminates after --cupti-timeout-s seconds.

NVML Sampling

Lightweight GPU telemetry via pynvml:

from pynvml import nvml...
util, mem_util = nvmlDeviceGetUtilizationRates(device)
mem_used = nvmlDeviceGetMemoryInfo(device).used
temp = nvmlDeviceGetTemperature(device, SENSOR_CORE)

Timing: Sub-millisecond. NVML util is windowed on some drivers.

Windows-Specific Considerations

Configuration & Customization

Adjusting Problem Sizes

To test different matrix sizes, edit the sweep configuration in experiment scripts or pass arguments:

python phase3_rollout_log.py --matrix-sizes 128 256 512 1024 --kernels gemm softmax

Modifying Knob Values

Block sizes: Edit environment/action_space.py

BLOCK_SIZES = [64, 128, 256, 512]  # Add or remove sizes

Register caps: Edit environment/action_space.py

REG_CAPS = [0, 16, 32, 48, 64]  # Add more granular caps

Tuning PPO Training

See Phase 3 documentation for detailed hyperparameter tuning. Start with defaults, then adjust learning_rate or n_epochs if convergence is slow.

Running on Different GPUs

Most changes are automatic (occupancy formulae adapt to sm_version). However, verify:

• GPU memory sufficient for batch sizes
• Compute capability ≥ 6.0 (modern NVIDIA)
• CUDA/Numba support for the GPU architecture

API Reference

KernelOptimizationEnv

from environment.kernel_env import KernelOptimizationEnv, EpisodeConfig

cfg = EpisodeConfig(
    kernel_name="gemm",          # or "random"
    matrix_size=512,
    max_steps=50,
    use_cupti=False,
    use_nvml=True,
    cupti_timeout_s=120
)

env = KernelOptimizationEnv(cfg)
obs, info = env.reset(seed=0)

for _ in range(10):
    action = env.action_space.sample()  # Random action
    obs, reward, terminated, truncated, info = env.step(action)
    if terminated or truncated:
        break

PPO Training

from stable_baselines3 import PPO

# Load environment
env = KernelOptimizationEnv(cfg)

# Create and train agent
model = PPO("MlpPolicy", env, learning_rate=3e-4, device="cuda")
model.learn(total_timesteps=50000, progress_bar=True)

# Save
model.save("models/ppo_trained")

# Load and use
loaded_model = PPO.load("models/ppo_trained")
obs, _ = env.reset()
action, _ = loaded_model.predict(obs, deterministic=True)
obs, reward, _, _, _ = env.step(action)

Metric Dictionaries

from profiling.cupti_collector import DEFAULT_NCU_METRICS

# Predefined CUPTI metrics
metrics = {
    "achieved_occupancy": DEFAULT_NCU_METRICS["achieved_occupancy"],
    "l2_hit_rate": DEFAULT_NCU_METRICS["l2_hit_rate"],
    "dram_bw_pct": DEFAULT_NCU_METRICS["dram_bw_pct"],
    "sm_active_pct": DEFAULT_NCU_METRICS["sm_active_pct"],
}

Troubleshooting

CUDA Errors

Profiling Errors

Training Issues

Windows-Specific

Help Understanding Experiment Results

This section embeds the full guide from help-understanding.md directly into the docs website so readers can understand the experiments, outputs, and how to interpret results without leaving this page.

Phase 5: BiLSTM Phase Detector

Purpose

Train a Bidirectional LSTM neural network that classifies GPU kernel execution into one of four phases based on temporal sequences of hardware performance counters. This answers: "What regime is the GPU operating in right now?"

The Four Execution Phases

Architecture

Input: (batch, T=20, 5 CUPTI counters)
  → BiLSTM (2 layers, hidden=64, bidirectional) → (batch, 128)
  → Phase Head: Linear(128→32) → ReLU → Linear(32→4) → Softmax
  → Uncertainty Head: Linear(128→16) → ReLU → Linear(16→1) → Sigmoid

Parameters: 142,021

How to Run

python training/train_phase_detector.py

Time: ~30 seconds | Output: results/models/phase_detector.pt, results/tables/phase5_eval.csv

Results

100% accuracy on synthetic roofline-labeled data. Real CUPTI traces expected: 85–95%.

Phase 6: GNN IR Encoder

Purpose

Encode the kernel's compiled PTX intermediate representation as a graph and use a Graph Convolutional Network (GCN) to produce a fixed-size structural embedding. This gives the RL agent kernel-aware context beyond runtime counters alone.

Pipeline

PTX source → Split into basic blocks (nodes) + control flow (edges)
  → 3 × GCNConv(10→64→64→64) → global_mean_pool
  → Linear(64→64) → concat(5 global features)
  → Output: 69-dim kernel structure embedding

Parameters: 18,014

Node Features (10-dim per basic block)

n_instructions, n_loads, n_stores, n_fma, n_add, n_mul, n_branches, n_sync, n_cvt, n_mov

Global Features (5-dim per kernel)

register_count, shared_memory, arithmetic_intensity, sync_count, instruction_count (all normalized)

Real Kernel Results

Key Files

• compiler/ir_extractor.py — PTX extraction + PyG graph builder
• models/gnn_encoder.py — GCN encoder + KernelStructureCache

Phase 7: RL vs Baselines Comparison

Purpose

The main evaluation experiment that validates the entire project. Compares three strategies head-to-head on each kernel × problem size combination.

1. PTXAS default — compiler's default configuration (block_size=256, reg_cap=0)
2. Random search — sample 100 random configurations, keep the best
3. Trained PPO agent — deterministic policy from Phase 4 training

How to Run

python experiments/phase7_rl_vs_baselines.py --model results/models/rtx3050_01.zip

Time: ~5–10 minutes | Output: results/tables/phase7_comparison.csv

Results (RTX 3050 Ti)

Key Findings

1. Real speedups over PTXAS defaults: Up to 1.58× on softmax — static heuristics leave significant performance on the table.
2. PPO consistency advantage: 60× lower variance than random search (±0.002ms vs ±0.523ms) — critical for production deployment.
3. Kernel-dependent optimization: Compute-bound kernels (GEMM) are well-served by defaults; memory-bound kernels (reduction, softmax) benefit most from adaptive tuning.
4. RL scales with action space: With 9 configs (3×3), random search explores exhaustively. More knobs will amplify the RL advantage.

Developer Information

Project Structure

├── kernels/              # CUDA kernels (GEMM, Reduction, Softmax)
├── profiling/            # Metric collection (CUPTI, NVML, timing)
├── compiler/             # Compilation control (PTXAS, IR extraction, GNN graph builder)
├── environment/          # Gymnasium RL environment (MDP, actions, states, reward)
├── models/               # Neural networks (PPO policy, BiLSTM phase detector, GCN encoder)
├── training/             # Training scripts (PPO, phase detector)
├── experiments/          # Phase experiments (Phase 0, 4, 7)
├── tests/                # Correctness and smoke tests
├── results/              # Output artifacts — CSV, models, logs (gitignored)
└── docs/                 # This documentation

Contributing

To contribute improvements:

1. Maintain all existing tests passing
2. Add tests for new functionality
3. Update documentation for user-facing changes
4. Follow PEP 8 style guide
5. Do not modify kernel implementations without explicit approval (ask maintainer)

Development Workflow

1. Create a feature branch
2. Run all tests: pytest tests/ -v
3. Run Phase 0/3 smoke tests to ensure no regressions
4. Submit pull request with detailed description

Performance Tips

• Use NVML-only mode during development (fast iteration)
• Profile with CUPTI sparingly (very slow)
• Cache compiled Numba kernels to speed up repeated runs
• Test changes on small problem sizes first

Developer Card

Acknowledgments

Built with NVIDIA Nsight Compute, Numba CUDA, Gymnasium, and Stable-Baselines3.