submission 37994

import torch
import time
import os
os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0"

from torch.utils.cpp_extension import load_inline

cuda_source = """
#define N_ROWS 8
#define THREADS 128

__global__ void __launch_bounds__(THREADS) sum_f32_to_f64_kernel(const float* x, double* out, size_t n) {
    __shared__ double sdata[3];
    const unsigned int tid = threadIdx.x;
    size_t idx = (blockIdx.x * blockDim.x + tid) * 4;
    const size_t row_size = gridDim.x * blockDim.x * 4;
    const unsigned int w_idx = tid % 32;

    double sum = 0.0;

    #pragma unroll
    for(int i = 0; i < N_ROWS; ++i){
        if(idx + 3 < n){
            const float4 data = __ldcs((float4*)(x + idx));
            sum += data.x;
            sum += data.y;
            sum += data.z;
            sum += data.w;

            idx += row_size;
        } else {
            for(; idx < n; ++idx)
                sum += __ldcs(x + idx);
        }   
    }

    sum += __shfl_down_sync(0xffffffff, sum, 16);
    sum += __shfl_down_sync(0xffffffff, sum, 8);
    sum += __shfl_down_sync(0xffffffff, sum, 4);
    sum += __shfl_down_sync(0xffffffff, sum, 2);
    sum += __shfl_down_sync(0xffffffff, sum, 1);

    if(w_idx == 0 and tid != 0){
        sdata[tid/32-1] = sum;
    }

    __syncthreads();

    if (tid == 0){
        sum += sdata[0];
        sum += sdata[1];
        sum += sdata[2];
        atomicAdd(out, sum);
    }
}

torch::Tensor sum_f32_to_f64(torch::Tensor x) {
    auto out = torch::zeros({1}, torch::dtype(torch::kFloat64).device(x.device()));

    int threads = THREADS * (4 * N_ROWS);
    int blocks = (x.numel() + threads - 1) / threads;

    sum_f32_to_f64_kernel<<<blocks, THREADS>>>(
        x.data_ptr<float>(), out.data_ptr<double>(), x.numel()
    );

    return out;
}
"""

cpp_source = """
torch::Tensor sum_f32_to_f64(torch::Tensor x);
"""

# Compile inline
module = load_inline(
    name="sum_f32_to_f64",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["sum_f32_to_f64"],
    verbose=False,
    with_cuda=True,
    extra_cuda_cflags=['-arch=compute_80', '-O3', '-arch=native']
)

from task import input_t, output_t

def custom_kernel(data: input_t) -> output_t:
    return module.sum_f32_to_f64(data[0])[0].to(torch.float32)