submission 607835

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

cuda_src = r'''
#include <torch/extension.h>
#include <cuda_runtime.h>

// Each thread processes 2 groups of 4 pixels (8 pixels total)
// Grid-stride loop to handle any size
__global__ void __launch_bounds__(256)
grayscale_kernel(const float* __restrict__ input,
                 float* __restrict__ output,
                 const int n_quads) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (tid < n_quads) {
        const float4* in4 = reinterpret_cast<const float4*>(input + tid * 12);
        float4 v0 = __ldg(in4);
        float4 v1 = __ldg(in4 + 1);
        float4 v2 = __ldg(in4 + 2);

        float4 out;
        out.x = __fmaf_rn(0.2989f, v0.x, __fmaf_rn(0.5870f, v0.y, 0.1140f * v0.z));
        out.y = __fmaf_rn(0.2989f, v0.w, __fmaf_rn(0.5870f, v1.x, 0.1140f * v1.y));
        out.z = __fmaf_rn(0.2989f, v1.z, __fmaf_rn(0.5870f, v1.w, 0.1140f * v2.x));
        out.w = __fmaf_rn(0.2989f, v2.y, __fmaf_rn(0.5870f, v2.z, 0.1140f * v2.w));

        reinterpret_cast<float4*>(output)[tid] = out;
    }
}

torch::Tensor launch_grayscale(torch::Tensor input, torch::Tensor output) {
    const int n_pixels = input.size(0) * input.size(1);
    const int n_quads = n_pixels >> 2;
    
    constexpr int threads = 256;
    const int blocks = (n_quads + 255) >> 8;
    
    grayscale_kernel<<<blocks, threads>>>(
        input.data_ptr<float>(),
        output.data_ptr<float>(),
        n_quads
    );
    
    return output;
}
'''

cpp_src = '''
torch::Tensor launch_grayscale(torch::Tensor input, torch::Tensor output);
'''

module = load_inline(
    name='grayscale_best_final',
    cpp_sources=cpp_src,
    cuda_sources=cuda_src,
    functions=['launch_grayscale'],
    verbose=False,
    extra_cuda_cflags=['-O3', '--use_fast_math'],
)

_launch = module.launch_grayscale

def custom_kernel(data: input_t) -> output_t:
    x, output = data
    return _launch(x, output)