submission 781388

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

convert_cuda_source = """
#define CEIL_DIV(a,b) (((a) + (b) - 1) / (b))
#define THREADS_PER_BLOCK 128

__global__ void grayscale(const float* __restrict__ rgb_image, float* __restrict__ grayscale_output) {
    int bX = blockIdx.x;
    int tX = threadIdx.x;

    int globalX = (bX * blockDim.x + tX);

    const float4 * inputFloat4 = reinterpret_cast<const float4*>(rgb_image);


    int baseIndex = globalX * 3; // 3 rgb values for each cell
    float4 l0 = inputFloat4[baseIndex + 0];
    float4 l1 = inputFloat4[baseIndex + 1];;
    float4 l2 = inputFloat4[baseIndex + 2];;


        
    float4 output;
    output.x = __fmaf_rn(0.299f, l0.x, __fmaf_rn(0.587f, l0.y, 0.114f * l0.z));
    output.y = __fmaf_rn(0.299f, l0.w, __fmaf_rn(0.587f, l1.x, 0.114f * l1.y));
    output.z = __fmaf_rn(0.299f, l1.z, __fmaf_rn(0.587f, l1.w, 0.114f * l2.x));
    output.w = __fmaf_rn(0.299f, l2.y, __fmaf_rn(0.587f, l2.z, 0.114f * l2.w));
        
    *reinterpret_cast<float4 *>(&grayscale_output[globalX*4]) = output;        
    
}


torch::Tensor convert_cuda(torch::Tensor input, torch::Tensor output) {
    TORCH_CHECK(input.device().is_cuda(), "Tensor input must be a CUDA tensor");
    TORCH_CHECK(output.device().is_cuda(), "Tensor output must be a CUDA tensor");
    
    auto sizes = input.sizes(); // retourne IntArrayRef
    int64_t height = sizes[0];
    int64_t width = sizes[1];

   int gDim = CEIL_DIV(height*width, THREADS_PER_BLOCK*4);

    int nbBlocks = gDim;
    grayscale<<<nbBlocks, THREADS_PER_BLOCK>>>(input.data_ptr<float>(), output.data_ptr<float>());


    

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        throw std::runtime_error(cudaGetErrorString(err));
    }

    return output;
}
"""

convert_cpp_source = """
#include <torch/extension.h>

torch::Tensor convert_cuda(torch::Tensor input, torch::Tensor output);
"""

convert_module = load_inline(
    name='convert_cuda',
    cpp_sources=convert_cpp_source,
    cuda_sources=convert_cuda_source,
    functions=['convert_cuda'],
    verbose=True,
)

def convert(input,output):
    if not input.is_cuda or not output.is_cuda:
        raise RuntimeError("Tensor must be on GPU")
    return convert_module.convert_cuda(input,output)

def custom_kernel(data: input_t) -> output_t:
    """
    Custom implementation of vector sum reduction using CUDA.
    Args:
        inputs: List of pairs of tensors [A, B] to be added.
    Returns:
        Tensor containing element-wise sum.
    """
    input, output = data
    assert input.is_cuda, "Input tensor must be on GPU"

    # Simply reuse the existing add function we already defined
    # This avoids the compilation issues with the inline kernel
    res =  convert(input, output)
    return res