Triton

import triton
import triton.language as tl
import torch

@triton.jit
def add_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    # Map program_id to a block of elements
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)

    # Create a mask to handle out-of-bounds accesses
    mask = offsets < n_elements

    # Load data from memory
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)

    # Perform addition
    output = x + y

    # Write back to memory
    tl.store(output_ptr + offsets, output, mask=mask)


def add(x: torch.Tensor, y: torch.Tensor):
    output = torch.empty_like(x)
    assert x.is_cuda and y.is_cuda and output.is_cuda
    n_elements = x.numel()

    # The block size is a compile-time constant, so we can't use `n_elements`
    # directly. Instead, we use a heuristic to choose a good block size.
    BLOCK_SIZE = 1024 # Or adjust based on your needs

    # Number of programs (blocks) to launch
    grid = lambda META: (triton.cdiv(n_elements, META['BLOCK_SIZE']),)

    # Launch the kernel
    add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=BLOCK_SIZE)
    return output

if __name__ == "__main__":
    # Example usage with PyTorch tensors
    size = 4096
    x = torch.randn(size, device='cuda')
    y = torch.randn(size, device='cuda')

    output_triton = add(x, y)
    output_torch = x + y

    print(f"Triton output matches PyTorch: {torch.allclose(output_triton, output_torch)}")