submission 66401

#!POPCORN leaderboard vectoradd_v2
import torch
import triton
import triton.language as tl


@triton.jit
def add_kernel(
    a_ptr,  # pointer to input A
    b_ptr,  # pointer to input B
    output_ptr,  # pointer to output
    n_elements,  # total number of elements
    BLOCK_SIZE: tl.constexpr,  # number of elements per block
):
    """Optimized element-wise addition kernel using Triton"""
    # Get program ID and compute element offset
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    
    # Create mask for boundary checking
    mask = offsets < n_elements
    
    # Load data with masking
    a = tl.load(a_ptr + offsets, mask=mask)
    b = tl.load(b_ptr + offsets, mask=mask)
    
    # Perform addition
    output = a + b
    
    # Store result with masking
    tl.store(output_ptr + offsets, output, mask=mask)


def custom_kernel(data):
    """
    High-performance vector addition using Triton CUDA kernel.
    Args:
        data: Tuple of tensors (A, B, output)
    Returns:
        output tensor with A + B
    """
    A, B, output = data
    
    # Get total number of elements
    n_elements = output.numel()
    
    # Choose optimal block size (tuned for modern GPUs)
    BLOCK_SIZE = 1024
    
    # Calculate grid size
    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
    
    # Launch kernel
    add_kernel[grid](
        A, B, output,
        n_elements,
        BLOCK_SIZE=BLOCK_SIZE,
    )
    
    return output