This paper introduces a practical guide to GPU programming using Mojo, a programming language that combines Python's syntax with system-level performance. Titled "Mojo 🔥 GPU Puzzles, Edition 1," the guide aims to democratize GPU programming by making it accessible to a broader audience, emphasizing a hands-on, puzzle-based learning approach.

The motivation for GPU programming is rooted in modern computing's reliance on parallel processing for tasks ranging from large language models and computer vision to scientific simulations and financial analysis. Traditional CPU performance is constrained by power and thermal limits, leading to the adoption of massively parallel architectures found in modern GPUs, such as the NVIDIA H100, which can execute thousands of threads concurrently.

Mojo is presented as an ideal language for this paradigm due to several key features:

• Python-style syntax: Lowering the barrier to entry for developers familiar with Python.
• Zero-cost abstractions: Compiling directly to machine code without performance overhead, even for system-level programming.
• Strong type system: Enabling early error detection at compile time.
• Tensor-native support: Designed with hardware optimization for tensor operations in mind.
• Direct hardware control: Allowing invocation of CPU and GPU intrinsics.
• Cross-hardware portability: Code can run efficiently on both CPUs and GPUs.
• Enhanced safety: Providing better memory safety compared to C/C++.

The core methodology of this guide is puzzle-based learning, which deviates from traditional approaches that front-load extensive theoretical knowledge. Instead, it encourages direct experimentation and discovery. Advantages of this approach include immediate feedback on GPU execution, progressive complexity building on previously learned concepts, a practical focus on real-world computational problems, development of debugging skills through systematic practice, and deeper understanding through active problem-solving rather than rote memorization. The guide draws inspiration from interactive NVIDIA GPU learning projects, re-implementing concepts with Mojo's unique abstractions and performance capabilities.

A fundamental aspect emphasized is the "GPU programming mindset," which involves a paradigm shift in how computations are viewed:

1. Sequential to Parallel: CPU programming typically uses sequential loops like $for i in range(data_size): result[i] = process(data[i])$. GPU programming transforms this by assigning thousands of threads to process data elements concurrently. Conceptually, each thread might execute $thread_id = get_global_id(); if thread_id < data_size: result[thread_id] = process(data[thread_id])$, where get_global_id() returns a unique identifier for the current thread within the global execution grid. This replaces explicit loops with massive parallel execution.
2. Fitting Computation Grid to Data: Effective GPU programming requires structuring computation as a grid of threads that maps efficiently onto the data. This involves understanding:
   • Threads: Individual processing units, each often responsible for a specific data element.
   • Blocks: Groups of threads that can share memory (shared memory) and synchronize execution within the group.
   • Grid: The overall hierarchical structure of blocks and threads that covers the entire computational problem.

The goal is to balance this thread configuration to maximize parallel efficiency while managing memory access patterns and synchronization requirements.

3. Data Movement vs. Computation Cost: A crucial insight is that in GPU programming, data movement costs often outweigh computation costs. Data transfer between CPU and GPU is slow. Moving data from global memory to shared memory is faster, and accessing data already in registers or shared memory is extremely fast. Therefore, optimization often focuses on minimizing data transfers and maximizing data locality, as the bottleneck is frequently data movement, not the arithmetic operations themselves.

The learning curriculum is structured into eleven parts, progressively covering topics from fundamental GPU concepts to advanced techniques:

• Part I (GPU Basics): Thread indexing, block configuration, memory access patterns, raw pointers vs. LayoutTensor, and shared memory.
• Part II (Debugging): GPU debuggers, sanitizers for memory errors and race conditions.
• Part III (GPU Algorithms): Parallel reduction, pooling, efficient convolution kernels, scan algorithms (prefix sum), and matrix multiplication optimization using tiling strategies.
• Part IV (MAX Graph Integration): Creating custom MAX graph operations and integrating GPU kernels with Python.
• Part V (PyTorch Integration): Connecting Mojo GPU kernels with PyTorch tensors, using CustomOpLibrary for tensor marshalling, torch.compile integration, kernel fusion, and custom backward passes.
• Part VI (Functional Patterns & Benchmarking): elementwise, tiled processing, vectorization, performance optimization tradeoffs, quantitative benchmarking, and understanding GPU threading vs. SIMD execution hierarchy.
• Part VII (Warp-level Programming): Warp basics, SIMT execution model, core warp operations (e.g., sum, shuffle_down, broadcast), and advanced patterns using shuffle_xor and prefix_sum.
• Part VIII (Block-level Programming): Block-level reductions (block.sum(), block.max()), block-level scan patterns, and efficient coordination within blocks using block.broadcast().
• Part IX (Advanced Memory System): Optimal memory coalescing, asynchronous memory operations to hide latency by overlapping computation and transfer, memory fences, synchronization primitives, prefetching, and cache optimization strategies.
• Part X (Performance Analysis & Optimization): GPU kernel profiling, occupancy and resource utilization optimization, and eliminating shared memory bank conflicts.
• Part XI (Advanced GPU Features): Tensor Core programming for AI workloads and cluster programming on modern GPUs.

The guide distinguishes itself by initially focusing on low-level memory manipulation to build a deep understanding of GPU memory patterns, then progressively transitioning to Mojo's higher-level LayoutTensor abstraction, providing both foundational knowledge and practical skills for modern tensor-based approaches.