submission 516101

#!POPCORN leaderboard grayscale_v2
#!POPCORN gpu A100

import torch
import triton
import triton.language as tl

@triton.jit
def grayscale_kernel(
    input_ptr,
    output_ptr,
    H,
    W,
    BLOCK: tl.constexpr
):
    pid = tl.program_id(0)
    pixel_idx = pid * BLOCK + tl.arange(0, BLOCK)
    mask = pixel_idx < H * W

    base = pixel_idx * 3
    R = tl.load(input_ptr + base + 0, mask=mask, other=0.0)
    G = tl.load(input_ptr + base + 1, mask=mask, other=0.0)
    B = tl.load(input_ptr + base + 2, mask=mask, other=0.0)

    gray = 0.2989 * R + 0.5870 * G + 0.1140 * B
    tl.store(output_ptr + pixel_idx, gray, mask=mask)

def custom_kernel(data):
    input_tensor, output_tensor = data
    H, W, _ = input_tensor.shape
    N = H * W

    BLOCK = 1024
    grid = (triton.cdiv(N, BLOCK),)

    grayscale_kernel[grid](
        input_tensor,
        output_tensor,
        H,
        W,
        BLOCK=BLOCK
    )
    return output_tensor