submission 614526

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>

__global__ void __launch_bounds__(1024, 1)
grayscale_kernel(const float* __restrict__ input,
                 float* __restrict__ output,
                 const int n_quads) {
    const int tid = blockIdx.x * 1024 + threadIdx.x;

    if (tid < n_quads) {
        const int base = tid * 12;
        const float4 v0 = __ldg((const float4*)(input + base));
        const float4 v1 = __ldg((const float4*)(input + base + 4));
        const float4 v2 = __ldg((const float4*)(input + base + 8));

        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));

        float* dst = (float*)output + tid * 4;
        asm volatile(
            "st.global.cs.v4.f32 [%0], {%1,%2,%3,%4};\n\t"
            :: "l"(dst), "f"(out.x), "f"(out.y), "f"(out.z), "f"(out.w)
            : "memory"
        );
    }
}

torch::Tensor launch_grayscale(torch::Tensor input, torch::Tensor output) {
    const int n_quads = (input.size(0) * input.size(1)) >> 2;
    grayscale_kernel<<<(n_quads + 1023) / 1024, 1024>>>(
        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_v7_cs',
    cpp_sources=cpp_src,
    cuda_sources=cuda_src,
    functions=['launch_grayscale'],
    verbose=False,
    extra_cuda_cflags=['-O3', '--use_fast_math', '-maxrregcount=28'],
)

_launch = module.launch_grayscale

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