submission 512272

#!POPCORN leaderboard grayscale_v2
#!POPCORN gpu A100

import torch
import triton
import triton.language as tl

from task import input_t, output_t


@triton.autotune(
    configs=[
        triton.Config({"BLOCK": 256}, num_warps=4, num_stages=2),
        triton.Config({"BLOCK": 512}, num_warps=4, num_stages=2),
        triton.Config({"BLOCK": 1024}, num_warps=8, num_stages=2),
        triton.Config({"BLOCK": 2048}, num_warps=8, num_stages=2),
    ],
    key=["n_pixels"],
)
@triton.jit
def _grayscale_kernel(
    x_ptr,
    y_ptr,
    n_pixels,
    BLOCK: tl.constexpr,
):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n_pixels

    base = offs * 3
    r = tl.load(x_ptr + base + 0, 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 = r * 0.2989 + g * 0.5870 + b * 0.1140
    tl.store(y_ptr + offs, y, mask=mask)


def _torch_fallback(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    y[...] = x[..., 0] * 0.2989 + x[..., 1] * 0.5870 + x[..., 2] * 0.1140
    return y


def custom_kernel(data: input_t) -> output_t:
    x, y = data

    if (
        not x.is_cuda
        or not y.is_cuda
        or x.dtype != torch.float32
        or y.dtype != torch.float32
        or x.ndim != 3
        or x.shape[-1] != 3
    ):
        return _torch_fallback(x, y)

    if not x.is_contiguous():
        x = x.contiguous()
    if not y.is_contiguous():
        y = y.contiguous()

    h, w, _ = x.shape
    n_pixels = h * w
    grid = lambda meta: (triton.cdiv(n_pixels, meta["BLOCK"]),)

    try:
        _grayscale_kernel[grid](x, y, n_pixels)
        return y
    except Exception:
        return _torch_fallback(x, y)