submission 545196

import os
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

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

@triton.jit
def grayscale_kernel(
    data_ptr, out_ptr, n_pixels,
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_pixels
    
    # data is (H, W, 3) = contiguous [R0, G0, B0, R1, G1, B1, ...]
    base = offsets * 3
    r = tl.load(data_ptr + base, mask=mask, other=0.0)
    g = tl.load(data_ptr + base + 1, mask=mask, other=0.0)
    b = tl.load(data_ptr + base + 2, mask=mask, other=0.0)
    
    gray = r * 0.2989 + g * 0.5870 + b * 0.1140
    tl.store(out_ptr + offsets, gray, mask=mask)

def custom_kernel(data: input_t) -> output_t:
    data, output = data
    h, w, _ = data.shape
    n_pixels = h * w
    BLOCK_SIZE = 2048
    grid = ((n_pixels + BLOCK_SIZE - 1) // BLOCK_SIZE,)
    grayscale_kernel[grid](data, output, n_pixels, BLOCK_SIZE=BLOCK_SIZE)
    return output