submission 66605

import torch, triton, triton.language as tl
from task import input_t, output_t

@triton.jit
def rgb2gray_kernel(
    x_ptr, y_ptr,
    H, W,
    stride_h, stride_w, stride_c,
    y_stride_h, y_stride_w,
    w_r, w_g, w_b,
    BLOCK_H: tl.constexpr,  # <-- constexpr tile sizes
    BLOCK_W: tl.constexpr
):
    pid_h = tl.program_id(0)
    pid_w = tl.program_id(1)

    h0 = pid_h * BLOCK_H
    w0 = pid_w * BLOCK_W

    hs = h0 + tl.arange(0, BLOCK_H)
    ws = w0 + tl.arange(0, BLOCK_W)

    Hs = hs[:, None]   # (BH, 1)
    Ws = ws[None, :]   # (1, BW)

    mask = (Hs < H) & (Ws < W)

    base = Hs * stride_h + Ws * stride_w  # offsets in elements

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

    #y = w_r * r + w_g * g + w_b * b
    y = tl.math.fma(w_b, b, tl.math.fma(w_g, g, w_r * r))

    tl.store(y_ptr + Hs * y_stride_h + Ws * y_stride_w, y, mask=mask)

def custom_kernel(data: input_t) -> output_t:
    x, out = data  # x: (H,W,3) float32 cuda contiguous; out: (H,W) float32 cuda
    H, W, C = x.shape
    assert C == 3
    # PyTorch gives strides in elements for contiguous NHWC: (W*3, 3, 1)
    s_h, s_w, s_c = x.stride()
    ys_h, ys_w = out.stride()

    # Tile/grid setup. Try (64, 64) first; (32,128) or (128,64) may be faster depending on size.
    BLOCK_H, BLOCK_W = 64, 64
    grid = (triton.cdiv(H, BLOCK_H), triton.cdiv(W, BLOCK_W))

    rgb2gray_kernel[grid](
        x, out,
        H, W,
        s_h, s_w, s_c,
        ys_h, ys_w,
        0.2989, 0.5870, 0.1140,
        BLOCK_H=BLOCK_H, BLOCK_W=BLOCK_W,   # <-- pass constexprs
        num_warps=4,                        # tune: 4 or 8
        num_stages=2                        # tune: 2 or 3
    )
    return out