submission 781371

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

__global__ void grayscale(const float* rgb_image, float* grayscale_output, size_t height, size_t width, int cellsPerThread) {
    int bX = blockIdx.x;
    int bY = blockIdx.y;
    int tX = threadIdx.x;
    int tY = threadIdx.y;

    int xCoord = bX * blockDim.x + tX;
    int yCoord = (bY * blockDim.y + tY) * cellsPerThread;


    for (int cell = 0; cell<cellsPerThread; cell+=4) {
        float4 l0 = *reinterpret_cast<const float4 *>(&rgb_image[xCoord * width * 3 + yCoord * 3 + 0*4]);
        float4 l1 = *reinterpret_cast<const float4 *>(&rgb_image[xCoord * width * 3 + yCoord * 3 + 1*4]);
        float4 l2 = *reinterpret_cast<const float4 *>(&rgb_image[xCoord * width * 3 + yCoord * 3 + 2*4]);

        float r0 = l0.x; float g0 = l0.y; float b0 = l0.z;
        float r1 = l0.w; float g1 = l1.x; float b1 = l1.y;
        float r2 = l1.z; float g2 = l1.w; float b2 = l2.x;
        float r3 = l2.y; float g3 = l2.z; float b3 = l2.w;
        
        float4 output = make_float4(
            0.299f * r0 + 0.587f * g0 + 0.114f * b0,
            0.299f * r1 + 0.587f * g1 + 0.114f * b1,
            0.299f * r2 + 0.587f * g2 + 0.114f * b2,
            0.299f * r3 + 0.587f * g3 + 0.114f * b3
        );

        
        *reinterpret_cast<float4 *>(&grayscale_output[xCoord * width + yCoord + cell]) = 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 blockSize = 16;
    int cellsPerThread = 4;
    int bDimX = CEIL_DIV(height, blockSize);
    int bDimY = CEIL_DIV(width, blockSize*cellsPerThread);
    dim3 threadsPerBlock(blockSize,blockSize);
    dim3 nbBlocks( bDimX, bDimY);

    grayscale<<<nbBlocks, threadsPerBlock>>>(input.data_ptr<float>(), output.data_ptr<float>(), height, width, cellsPerThread);


    

    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