submission 43581

import torch
import triton
import triton.language as tl
#!POPCORN leaderboard identity_py
from task import input_t, output_t

@triton.jit
def kernel(inp, out, numel, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    offs = tl.arange(0, BLOCK_SIZE) + pid * BLOCK_SIZE
    red = tl.load(inp + offs * 3, mask=(offs * 3< numel))
    green = tl.load(inp + offs * 3 + 1, mask=(offs * 3 + 1 < numel))
    blue = tl.load(inp + offs * 3 + 2, mask=(offs * 3 + 2 < numel))
    res = 0.2989 * red + 0.5870 * green + 0.1140 * blue
    tl.store(out + offs, res, mask=(offs * 3 < numel))

# User kernel implementation.
def custom_kernel(input: input_t) -> output_t:
    data, output = input
    numel = data.numel()
    BLOCK_SIZE = 128
    grid = (triton.cdiv(numel, 3 * BLOCK_SIZE),)
    kernel[grid](data, output, numel, BLOCK_SIZE)
    return output