submission 695980

import triton
import triton.language as tl
import torch


@triton.jit
def add_kernel(
    A_ptr, B_ptr, C_ptr,
    n_elements,
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(0)

    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements

    tl.multiple_of(offsets, 16)

    a = tl.load(A_ptr + offsets, mask=mask, eviction_policy='evict_last')
    b = tl.load(B_ptr + offsets, mask=mask, eviction_policy='evict_last')

    c = a + b

    tl.store(C_ptr + offsets, c, mask=mask)


def custom_kernel(data):
    A, B, C = data

    n_elements = A.numel()

    BLOCK_SIZE = 8192

    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)

    add_kernel[grid](
        A, B, C,
        n_elements,
        BLOCK_SIZE=BLOCK_SIZE,
        num_warps=4,
        num_stages=2,
    )

    return C