Developing expertise in high-performance computing optimization requires systematic practice with modern programming models. Triton has emerged as a compelling framework for GPU programming, offering a balance between low-level control and high-level abstractions. This article presents an analysis of the first seven puzzles from the Triton-Puzzles-Lite repository, examining key concepts and implementation strategies for efficient GPU kernel development.

Exercise Repository

The complete exercise set is available at:
https://github.com/SiriusNEO/Triton-Puzzles-Lite

Exercise 1: Constant Addition

Objective

Add a constant value of 10 to every element in an input vector.

Analysis

This foundational exercise establishes essential patterns for pointer arithmetic and data movement in Triton. The kernel demonstrates the core workflow: generating indices, loading data from memory, performing computations, and storing results back to global memory.

@triton.jit
def constant_add_kernel(input_ptr, output_ptr, vector_size, tile_size: tl.constexpr):
    # Generate tile indices
    tile_indices = tl.arange(0, tile_size)
    
    # Load source data using pointer offset
    source_data = tl.load(input_ptr + tile_indices)
    
    # Perform element-wise addition
    result = source_data + 10.0
    
    # Store computed results
    tl.store(output_ptr + tile_indices, result)

Key Concepts

• Index generation using tl.arange creates consecutive integer sequences
• Pointer arithmetic enables efficient memory access patterns
• The decorated function compiles to optimized GPU code

Exercise 2: Block-Based Constant Addition

Objective

Extend the previous exercise to handle vectors larger than the block size using blocked processing.

Analysis

When working with large datasets, processing data in blocks becomes essential. This exercise introduces program identity tracking and boundary masking to insure safe memory access across block boundaries.

@triton.jit
def blocked_constant_add_kernel(input_ptr, output_ptr, total_length, block_size: tl.constexpr):
    # Determine block offset using program identifier
    block_offset = tl.program_id(0) * block_size
    tile_indices = block_offset + tl.arange(0, block_size)
    
    # Create validity mask for boundary conditions
    valid_mask = tile_indices < total_length
    
    # Load data with masking to prevent out-of-bounds access
    source_data = tl.load(input_ptr + tile_indices, mask=valid_mask)
    
    # Compute results
    result = source_data + 10.0
    
    # Store with appropriate masking
    tl.store(output_ptr + tile_indices, result, mask=valid_mask)

Key Concepts

• tl.program_id identifies the current block being procesed
• Masking ensures memory operations respect array boundaries
• Block processing enables handling arbitrarily large datasets

Exercise 3: Outer Vector Addition

Objective

Compute the outer product sum between two vectors, broadcasting across dimensions.

Analysis

This exercise explores two-dimensional computation patterns. The outer product creates a matrix where each element represents the sum of corresponding elements from the input vectors. Broadcasting enables efficient element-wise operations across different dimensions.

@triton.jit
def outer_vector_add_kernel(vec_a_ptr, vec_b_ptr, result_ptr, dim_a, dim_b, tile_a: tl.constexpr, tile_b: tl.constexpr):
    # Create two-dimensional index grids
    row_indices = tl.arange(0, tile_a)[:, None]
    col_indices = tl.arange(0, tile_b)[None, :]
    
    # Load corresponding vector elements
    vector_a = tl.load(vec_a_ptr + row_indices)
    vector_b = tl.load(vec_b_ptr + col_indices)
    
    # Broadcast addition across dimensions
    computed_result = vector_a + vector_b
    
    # Store with linearized two-dimensional indices
    linear_index = row_indices * dim_b + col_indices
    tl.store(result_ptr + linear_index, computed_result)

Key Concepts

• Two-dimensional index structures enable matrix generation
• Broadcasting simplifies element-wise operations across tensor dimensions
• Two-dimensional indices require linearization for storage

Exercise 4: Block-Based Outer Vector Addition

Objective

Apply blocked processing to outer vector addition for large-scale computations.

Analysis

Combining blocked processing with two-dimensional computation maximizes GPU utilization. Multiple program identifiers track position in both dimensions, enabling parallel processing of matrix blocks.

@triton.jit
def blocked_outer_add_kernel(vec_a_ptr, vec_b_ptr, result_ptr, dim_a, dim_b, block_a: tl.constexpr, block_b: tl.constexpr):
    # Extract block coordinates
    block_row = tl.program_id(0)
    block_col = tl.program_id(1)
    
    # Calculate offsets within each dimension
    row_offset = block_row * block_a + tl.arange(0, block_a)
    col_offset = block_col * block_b + tl.arange(0, block_b)
    
    # Generate validity masks for boundary handling
    row_valid = row_offset < dim_a
    col_valid = col_offset < dim_b
    
    # Load blocked data segments
    segment_a = tl.load(vec_a_ptr + row_offset, mask=row_valid)
    segment_b = tl.load(vec_b_ptr + col_offset, mask=col_valid)
    
    # Compute outer product sum
    product_matrix = segment_a + segment_b
    
    # Store with combined masking
    result_linear_index = row_offset * dim_b + col_offset
    combined_mask = row_valid & col_valid
    tl.store(result_ptr + result_linear_index, product_matrix, mask=combined_mask)

Key Concepts

• Multi-dimensional program identifiers enable sophisticated parallel patterns
• Combined masks handle irregular-sized dimensions
• Blocked operations improve memory access efficiency

Exercise 5: Fused Outer Multiplication with Activation

Objective

Compute the outer product of two vectors and apply the ReLU activation function.

Analysis

This exercise introduces non-linear activation functions into the blocked outer product pattern. The fusion of multiplication and activation demonstrates Triton's capability for expressing complex operations efficiently.

@triton.jit
def fused_mul_relu_kernel(vec_x_ptr, vec_y_ptr, output_ptr, size_x, size_y, tile_x: tl.constexpr, tile_y: tl.constexpr):
    # Determine current block position
    block_idx_x = tl.program_id(0)
    block_idx_y = tl.program_id(1)
    
    # Compute actual offsets
    offset_x = block_idx_x * tile_x + tl.arange(0, tile_x)
    offset_y = block_idx_y * tile_y + tl.arange(0, tile_y)
    
    # Generate boundary masks
    valid_x = offset_x < size_x
    valid_y = offset_y < size_y
    
    # Load data with proper masking
    data_x = tl.load(vec_x_ptr + offset_x, mask=valid_x)
    data_y = tl.load(vec_y_ptr + offset_y, mask=valid_y)
    
    # Perform element-wise multiplication
    product = data_x * data_y
    
    # Apply ReLU activation using conditional operation
    activated = tl.where(product > 0, product, 0.0)
    
    # Store results with combined validity check
    linear_pos = offset_x * size_y + offset_y
    storage_mask = valid_x & valid_y
    tl.store(output_ptr + linear_pos, activated, mask=storage_mask)

Key Concepts

• Element-wise multiplication forms the basis of outer product computation
• tl.where implements conditional operations efficiently
• Kernel fusion reduces memory traffic and improves performance

Exercise 6: Backward Pass for Fused Multiplication

Objective

Implement the gradient computation for the outer product with ReLU activation.

Analysis

Gradient computation requires chaining derivatives through the forward operations. This exercise demonstrates how Triton's primitives support automatic differentiation patterns, with gradient flow requiring careful consideration of intermediate values.

@triton.jit
def mul_relu_gradient_kernel(x_ptr, y_ptr, grad_output_ptr, grad_x_ptr, dim_x, dim_y, tile_x: tl.constexpr, tile_y: tl.constexpr):
    # Identify current processing block
    block_x = tl.program_id(0)
    block_y = tl.program_id(1)
    
    # Calculate tile offsets
    off_x = block_x * tile_x + tl.arange(0, tile_x)
    off_y = block_y * tile_y + tl.arange(0, tile_y)
    
    # Create boundary validity checks
    bounds_x = off_x < dim_x
    bounds_y = off_y < dim_y
    
    # Retrieve forward pass intermediate values
    input_x = tl.load(x_ptr + off_x, mask=bounds_x)
    input_y = tl.load(y_ptr + off_y, mask=bounds_y)
    
    # Load gradient with respect to output
    upstream_grad = tl.load(grad_output_ptr + off_x * dim_y + off_y, mask=bounds_x & bounds_y)
    
    # Compute gradient with respect to x using chain rule
    grad_x = upstream_grad * input_y
    
    # Store computed gradients
    tl.store(grad_x_ptr + off_x * dim_y + off_y, grad_x, mask=bounds_x & bounds_y)

Key Concepts

• Gradient computation relies on intermediate values from forward pass
• Chain rule application determines gradient distribution
• Consistent masking ensures correctness at boundaries

Exercise 7: Batched Vector Summation

Objective

Compute sum reductions across batched vector data.

Analysis

Reduction operations form the foundation of many machine learning workloads. This exercise demonstrates efficient summation across specific tensor dimensions while maintaining data independence between blocks.

@triton.jit
def batch_sum_kernel(data_ptr, output_ptr, outer_dim, inner_dim, stride, tile_size: tl.constexpr, inner_tile: tl.constexpr):
    # Determine processing block
    block_id = tl.program_id(0)
    
    # Calculate starting position for this block
    base_position = block_id * tile_size
    positions = base_position + tl.arange(0, tile_size)
    
    # Validate against outer dimension
    validity_mask = positions < outer_dim
    
    # Load batched data with appropriate stride
    batch_data = tl.load(
        data_ptr + positions * stride + tl.arange(0, inner_tile),
        mask=validity_mask
    )
    
    # Perform reduction along the inner dimension
    reduced = tl.sum(batch_data, 1)
    
    # Store results with boundary masking
    tl.store(output_ptr + positions, reduced, mask=validity_mask)

Key Concepts

• Reduction operations require specifying the axis for aggregation
• Strided access patterns enable flexible batch processing
• Block independence simplifies parallel implementation

Practical Insights

The mask operations throughout these exercises represent a critical pattern for writing robust GPU kernels. Proper boundary handling prevents memory access violations and ensures correctness across varying input sizes. These foundational exercises establish patterns applicable to more complex kernel implementations in production systems.