submission 617581

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>

// Branchless version: for power-of-2 sizes that are multiples of 4,
// n_quads is exactly divisible by 1024 so we skip the branch
__global__ void __launch_bounds__(1024, 1)
grayscale_kernel_exact(const float* __restrict__ input,
                       float* __restrict__ output) {
    const int tid = blockIdx.x * 1024 + threadIdx.x;
    const int base = tid * 12;

    const float2 a = __ldg((const float2*)(input + base));
    const float2 b = __ldg((const float2*)(input + base + 2));
    const float2 c = __ldg((const float2*)(input + base + 4));
    const float2 d = __ldg((const float2*)(input + base + 6));
    const float2 e = __ldg((const float2*)(input + base + 8));
    const float2 f = __ldg((const float2*)(input + base + 10));

    float4 out;
    out.x = __fmaf_rn(0.2989f, a.x, __fmaf_rn(0.5870f, a.y, 0.1140f * b.x));
    out.y = __fmaf_rn(0.2989f, b.y, __fmaf_rn(0.5870f, c.x, 0.1140f * c.y));
    out.z = __fmaf_rn(0.2989f, d.x, __fmaf_rn(0.5870f, d.y, 0.1140f * e.x));
    out.w = __fmaf_rn(0.2989f, e.y, __fmaf_rn(0.5870f, f.x, 0.1140f * f.y));

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

// Guarded version for non-exact sizes
__global__ void __launch_bounds__(1024, 1)
grayscale_kernel_guarded(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 float2 a = __ldg((const float2*)(input + base));
        const float2 b = __ldg((const float2*)(input + base + 2));
        const float2 c = __ldg((const float2*)(input + base + 4));
        const float2 d = __ldg((const float2*)(input + base + 6));
        const float2 e = __ldg((const float2*)(input + base + 8));
        const float2 f = __ldg((const float2*)(input + base + 10));

        float4 out;
        out.x = __fmaf_rn(0.2989f, a.x, __fmaf_rn(0.5870f, a.y, 0.1140f * b.x));
        out.y = __fmaf_rn(0.2989f, b.y, __fmaf_rn(0.5870f, c.x, 0.1140f * c.y));
        out.z = __fmaf_rn(0.2989f, d.x, __fmaf_rn(0.5870f, d.y, 0.1140f * e.x));
        out.w = __fmaf_rn(0.2989f, e.y, __fmaf_rn(0.5870f, f.x, 0.1140f * f.y));

        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;
    const int n_blocks = (n_quads + 1023) / 1024;

    // If exact multiple of 1024, use branchless kernel
    if ((n_quads & 1023) == 0) {
        grayscale_kernel_exact<<<n_blocks, 1024>>>(
            input.data_ptr<float>(),
            output.data_ptr<float>()
        );
    } else {
        grayscale_kernel_guarded<<<n_blocks, 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_v10',
    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)