submission 68709

import torch
import triton
import triton.language as tl

@triton.jit
def rgb2gray_fast(
    inp_ptr, out_ptr,
    BLOCK_SIZE: tl.constexpr,
):
    pid  = tl.program_id(0)
    i    = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)   # (B,)
    base = i + (i << 1)

    r = tl.load(inp_ptr + base + 0, cache_modifier=".ca")
    g = tl.load(inp_ptr + base + 1, cache_modifier=".ca")
    b = tl.load(inp_ptr + base + 2, cache_modifier=".ca")
    # gray = tl.fma(g, 0.5870, tl.fma(r, 0.2989, b * 0.1140))
    gray = tl.fma(r, 0.2989, tl.fma(b, 0.1140, g * 0.5870))
    # gray = tl.fma(b, 0.1140, tl.fma(g, 0.5870, r * 0.2989))
    tl.store(out_ptr + i, gray)

def custom_kernel(data: torch.Tensor):
    """
    Input:  (H, W, 3) contiguous, dtype=float32  (元コードと同じ前提)
    Output: (H, W)   dtype=float32
    """
    input, output = data

    H, W, _ = input.shape
    n_pixels = H * W

    BLOCK = 8192
    grid = (triton.cdiv(n_pixels, BLOCK), )

    rgb2gray_fast[grid](
        input, output,
        BLOCK_SIZE=BLOCK,
        num_warps=32,
        num_stages=2,
    )
    return output