submission 66559

import torch
import triton
import triton.language as tl

from task import input_t, output_t

DEVICE = triton.runtime.driver.active.get_active_torch_device()

# ---- Autotuned, vectorized kernel ------------------------------------------
@triton.autotune(
    configs=[
        # Good starting points across Ampere/Ada/Hopper; feel free to extend.
        triton.Config({'BLOCK_SIZE': 256,  'VEC': 8}, num_warps=4,  num_stages=2),
        triton.Config({'BLOCK_SIZE': 512,  'VEC': 8}, num_warps=8,  num_stages=2),
        triton.Config({'BLOCK_SIZE': 1024, 'VEC': 4}, num_warps=8,  num_stages=2),
        triton.Config({'BLOCK_SIZE': 2048, 'VEC': 2}, num_warps=8,  num_stages=2),
        triton.Config({'BLOCK_SIZE': 4096, 'VEC': 1}, num_warps=8,  num_stages=2),
    ],
    key=['n_elements'],  # autotune choice depends mainly on problem size
)
@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n_elements,
               BLOCK_SIZE: tl.constexpr, VEC: tl.constexpr):
    pid = tl.program_id(axis=0)

    # vectorized block start; each program handles BLOCK_SIZE * VEC elements
    block_start = pid * BLOCK_SIZE * VEC

    # 2D tile: [BLOCK_SIZE, VEC] so compiler emits wider loads/stores
    offs = block_start + tl.arange(0, BLOCK_SIZE)[:, None] * VEC + tl.arange(0, VEC)[None, :]
    mask = offs < n_elements

    # Optional: tell compiler about alignment/contiguity (safe for contiguous tensors)
    tl.multiple_of(block_start, VEC)

    x = tl.load(x_ptr + offs, mask=mask, other=0.0)
    y = tl.load(y_ptr + offs, mask=mask, other=0.0)
    z = x + y
    tl.store(out_ptr + offs, z, mask=mask)


def add_into(x: torch.Tensor, y: torch.Tensor, out: torch.Tensor):
    # Preconditions
    assert x.device == DEVICE and y.device == DEVICE and out.device == DEVICE
    assert x.is_contiguous() and y.is_contiguous() and out.is_contiguous()
    assert x.shape == y.shape == out.shape
    # (Your task already restricts to fp16, which is great for bandwidth.)
    assert x.dtype == torch.float16 and y.dtype == torch.float16 and out.dtype == torch.float16

    n_elements = out.numel()
    # Each program covers BLOCK_SIZE * VEC elements (meta values chosen by autotune).
    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE'] * meta['VEC']),)
    add_kernel[grid](x, y, out, n_elements)  # meta-params picked by autotune


def custom_kernel(data: input_t) -> output_t:
    """
    Custom implementation of vector addition using Triton (optimized).
    Args:
        inputs: triple (A, B, output)
    Returns:
        output (element-wise sum)
    """
    A, B, output = data
    assert A.is_cuda and B.is_cuda and output.is_cuda, "Tensors must be on GPU"
    assert A.shape == B.shape == output.shape, "All shapes must match"
    assert A.dtype == torch.float16 and B.dtype == torch.float16 and output.dtype == torch.float16, \
        "All tensors must be float16"
    add_into(A, B, output)
    return output