submission 66945

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,
    WR, WG, WB,                       # weights as scalars
    BLOCK_H: tl.constexpr,
    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)          # (BH,)
    ws = w0 + tl.arange(0, BLOCK_W)          # (BW,)
    Hs = hs[:, None]                          # (BH,1)
    Ws = ws[None, :]                          # (1,BW)

    mask = (Hs < H) & (Ws < W)
    base = Hs * stride_h + Ws * stride_w      # (BH,BW)

    # Vectorized 4-lane per-pixel load (last lane is dummy)
    ch = tl.arange(0, 4)                      # (4,)
    ptrs = x_ptr + base[:, :, None] + ch[None, None, :] * stride_c  # (BH,BW,4)

    rgb4 = tl.load(
        ptrs,
        mask=mask[:, :, None] & (ch[None, None, :] < 3),
        other=0.0,
        cache_modifier=".ca",
    )  # (BH,BW,4) = [R,G,B,0]

    # Build weights per lane without slicing
    # w_lane[k] = [WR, WG, WB, 0][k]
    w_lane = (tl.where(ch == 0, WR, 0.0)
            + tl.where(ch == 1, WG, 0.0)
            + tl.where(ch == 2, WB, 0.0))      # shape (4,)
    w_broadcast = w_lane[None, None, :]        # (1,1,4)

    # Weighted sum across channel axis
    y = tl.sum(rgb4 * w_broadcast, axis=2)     # (BH,BW)

    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
    H, W, C = x.shape
    s_h, s_w, s_c = x.stride()
    ys_h, ys_w = out.stride()

    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,          # WR, WG, WB
        BLOCK_H=BLOCK_H, BLOCK_W=BLOCK_W,
        num_warps=8, num_stages=2
    )
    return out