submission 688313

from task import input_t, output_t
import torch
from torch.utils.cpp_extension import load_inline

cuda_source = r"""
#include <torch/extension.h>

__global__ void grayscale_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    const int total_pixels
) {
    const float wr = 0.2989f;
    const float wg = 0.5870f;
    const float wb = 0.1140f;

    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;

    int n4 = total_pixels >> 2;
    for (int i = idx; i < n4; i += stride) {
        const float4* p = reinterpret_cast<const float4*>(input + i * 12);
        float4 a = __ldg(p);
        float4 b = __ldg(p + 1);
        float4 c = __ldg(p + 2);

        float4 o;
        o.x = __fmaf_rn(wr, a.x, __fmaf_rn(wg, a.y, wb * a.z));
        o.y = __fmaf_rn(wr, a.w, __fmaf_rn(wg, b.x, wb * b.y));
        o.z = __fmaf_rn(wr, b.z, __fmaf_rn(wg, b.w, wb * c.x));
        o.w = __fmaf_rn(wr, c.y, __fmaf_rn(wg, c.z, wb * c.w));

        reinterpret_cast<float4*>(output)[i] = o;
    }
}

torch::Tensor grayscale_cuda(torch::Tensor input, torch::Tensor output) {
    int total_pixels = input.size(0) * input.size(1);
    int n4 = total_pixels >> 2;

    const int threads = 256;
    int blocks = min((n4 + threads - 1) / threads, 65535);
    if (blocks < 1) blocks = 1;

    grayscale_kernel<<<blocks, threads>>>(
        input.data_ptr<float>(),
        output.data_ptr<float>(),
        total_pixels
    );

    return output;
}
"""

cpp_source = "torch::Tensor grayscale_cuda(torch::Tensor input, torch::Tensor output);"

module = load_inline(
    name="grayscale_fast",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["grayscale_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)

def custom_kernel(data: input_t) -> output_t:
    data, output = data
    module.grayscale_cuda(data, output)
    return output