submission 779876

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

cuda_source = """
#include <cuda_runtime.h>

__global__ void grayscale_kernel(
    const float* __restrict__ data,
    float* __restrict__ output
) {
    // each thread handles 4 pixels = 12 floats = 3 x float4 loads
    int base = (blockIdx.x * blockDim.x + threadIdx.x) * 4;

    float4 a = __ldg((const float4*)&data[base * 3 + 0]); // R0 G0 B0 R1
    float4 b = __ldg((const float4*)&data[base * 3 + 4]); // G1 B1 R2 G2
    float4 c = __ldg((const float4*)&data[base * 3 + 8]); // B2 R3 G3 B3

    float v0 = a.x * 0.2989f + a.y * 0.5870f + a.z * 0.1140f;
    float v1 = a.w * 0.2989f + b.x * 0.5870f + b.y * 0.1140f;
    float v2 = b.z * 0.2989f + b.w * 0.5870f + c.x * 0.1140f;
    float v3 = c.y * 0.2989f + c.z * 0.5870f + c.w * 0.1140f;

    // single 128-bit write-through store for all 4 output pixels
    asm("st.global.wt.v4.f32 [%0], {%1, %2, %3, %4};"
        :: "l"(&output[base]), "f"(v0), "f"(v1), "f"(v2), "f"(v3));
}

torch::Tensor grayscale_cuda(torch::Tensor data, torch::Tensor output) {
    int n_pixels = output.numel();
    const int BLOCK_SIZE = 256;
    int grid = n_pixels / (BLOCK_SIZE * 4);
    grayscale_kernel<<<grid, BLOCK_SIZE>>>(
        data.data_ptr<float>(),
        output.data_ptr<float>()
    );
    return output;
}
"""

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

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


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