submission 490594

import triton
import triton.language as tl
import torch


@triton.jit
def _vector_add_kernel(
    ptr_a,      # Pointer to input tensor A
    ptr_b,      # Pointer to input tensor B
    ptr_out,    # Pointer to output tensor C
    n_elements, # Total number of elements
    BLOCK_SIZE: tl.constexpr,  # Block size for parallelization
):
    """
    Triton kernel for element-wise addition of two tensors.
    
    Fused operation: C = A + B
    
    Each program instance processes BLOCK_SIZE elements.
    """
    # Get the program ID (which block we're processing)
    pid = tl.program_id(axis=0)
    
    # Calculate the starting offset for this block
    block_start = pid * BLOCK_SIZE
    
    # Generate offsets for elements within this block
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    
    # Create mask for boundary handling (handles non-power-of-2 sizes)
    mask = offsets < n_elements
    
    # Load elements from tensor A with masking
    a = tl.load(ptr_a + offsets, mask=mask, other=0.0)
    
    # Load elements from tensor B with masking
    b = tl.load(ptr_b + offsets, mask=mask, other=0.0)
    
    # Perform element-wise addition using Triton operations
    result = a + b
    
    # Store the result to output tensor with masking
    tl.store(ptr_out + offsets, result, mask=mask)


def kernel_function(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor) -> torch.Tensor:
    """
    Wrapper function for float16 vector addition kernel.
    
    Performs element-wise addition: C = A + B
    
    This is a single fused operation - no decomposition needed as it's
    already an atomic elementwise operation.
    
    Args:
        A: Input tensor of shape (N, N) and dtype float16
        B: Input tensor of shape (N, N) and dtype float16
        C: Output tensor of shape (N, N) and dtype float16 (pre-allocated)
    
    Returns:
        C: The output tensor containing A + B
    """
    # Validate inputs
    assert A.is_cuda and B.is_cuda and C.is_cuda, "All tensors must be on CUDA"
    assert A.dtype == torch.float16 and B.dtype == torch.float16, "Inputs must be float16"
    assert C.dtype == torch.float16, "Output must be float16"
    assert A.shape == B.shape == C.shape, "All tensors must have the same shape"
    assert A.is_contiguous() and B.is_contiguous() and C.is_contiguous(), "Tensors must be contiguous"
    
    # Calculate total number of elements
    n_elements = A.numel()
    
    # Choose block size (power of 2 for efficiency)
    BLOCK_SIZE = 1024
    
    # Calculate grid dimensions (number of blocks needed)
    grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
    
    # Launch the Triton kernel
    _vector_add_kernel[grid](
        A,          # Input tensor A
        B,          # Input tensor B
        C,          # Output tensor C
        n_elements, # Total elements to process
        BLOCK_SIZE, # Block size (compile-time constant)
    )
    
    return C

import inspect

def custom_kernel(input):
    sig = inspect.signature(kernel_function)
    num_params = len(sig.parameters)

    if len(input) == num_params:
        return kernel_function(*input)
    return kernel_function(input)


# Ensure deterministic cuBLAS.
import os
if os.environ.get("CUBLAS_WORKSPACE_CONFIG", "") not in (":4096:8", ":16:8"):
    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"