submission 571245

import triton
import triton.language as tl

from task import input_t, output_t


@triton.autotune(
    configs=[
        triton.Config({"BLOCK_SIZE": block_size}, num_warps=num_warps, num_stages=1)
        for block_size in (256, 512, 1024, 2048)
        for num_warps in (4, 8)
    ],
    key=["n"],
)
@triton.jit
def _rgb_to_grayscale_kernel(
    x_ptr,
    out_ptr,
    n,
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(0)
    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offs < n

    base = offs * 3
    r = tl.load(x_ptr + base, mask=mask, other=0.0)
    g = tl.load(x_ptr + base + 1, mask=mask, other=0.0)
    b = tl.load(x_ptr + base + 2, mask=mask, other=0.0)

    y = 0.2989 * r + 0.5870 * g + 0.1140 * b
    tl.store(out_ptr + offs, y, mask=mask)


# Keep a module-local alias to minimize attribute lookups in the timed path.
_launch_kernel = _rgb_to_grayscale_kernel
_cdiv = triton.cdiv


def custom_kernel(data: input_t) -> output_t:
    x, output = data
    n = output.numel()
    grid = lambda meta: (_cdiv(n, meta["BLOCK_SIZE"]),)
    _launch_kernel[grid](x, output, n)
    return output