submission 68131

#!POPCORN leaderboard vectorsum_v2
#!POPCORN gpus A100 H100 B200 L4
# @SUBMISSION@
# https://github.com/gpu-mode/profiling-cuda-in-torch/blob/main/load_inline.py
from task import input_t, output_t
import torch
from torch.utils.cpp_extension import load_inline

cuda_source = """
#include <pybind11/pybind11.h>
//#include <pybind11/tuple.h>
#include <torch/extension.h>

template <typename T>
__device__ T warpReduction(T threadSum, uint mask = 0xffffffff) {
    T warpSum {threadSum};
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        warpSum += __shfl_down_sync(mask, warpSum, offset);
    }
    return warpSum;
}

template <typename T, typename T2, uint blockSize>
__global__ void reduction_basic(const T* __restrict__ data, T2* __restrict__ out, size_t size) {
    uint index = blockIdx.x * blockDim.x + threadIdx.x;
    const uint stride = gridDim.x * blockDim.x;

    constexpr uint blockWarpCount = (blockSize + 31) / 32;
    constexpr uint blockReductionThreadMask = (1 << blockWarpCount) - 1;
    __shared__ T warpSums[blockWarpCount];

    T threadSum {};
    while (index < size) {
        threadSum += data[index];
        index += stride;
    }

    T warpSum = warpReduction(threadSum);

    //first thread of each warp writes the warp's sum to shared
    uint warpNumber = threadIdx.x >> 5;
    if (!(threadIdx.x & 0b11111)) {
        warpSums[warpNumber] = warpSum;
    }

    __syncthreads();

    //first warp reduces shared sums
    if (threadIdx.x < blockWarpCount) {
        T blockSum = warpReduction(warpSums[threadIdx.x], blockReductionThreadMask);
        //first thread writes result to global
        if (!threadIdx.x) {
            atomicAdd(out, (T2)blockSum);
        }
    }
}




template <typename T, typename T2, uint blockSize, uint pipelineDepth = 1>
__global__ void reduction_vectorized(const T* __restrict__ data, T2* __restrict__ out, size_t size) {
    uint index = blockIdx.x * blockDim.x + threadIdx.x;
    const uint stride = gridDim.x * blockDim.x;

    constexpr uint blockWarpCount = (blockSize + 31) / 32;
    constexpr uint blockReductionThreadMask = (1 << blockWarpCount) - 1;
    __shared__ T warpSums[blockWarpCount];

    constexpr int fPerF4 = sizeof(float4) / sizeof(float);
	constexpr int floatsPerChunk = pipelineDepth * fPerF4;
    union localBuffer {
		float4 f4[pipelineDepth];
		float f[pipelineDepth * fPerF4];
	};

    localBuffer localX;
	const int vectorizedIterationCount = size / floatsPerChunk;

    T threadSum {};
    while (index < vectorizedIterationCount) {
        #pragma unroll
		for (int i = 0; i < pipelineDepth; ++i) {
            localX.f4[i] = reinterpret_cast<const float4*>(data)[index * pipelineDepth + i];
		}
		#pragma unroll
		for (int i = 0; i < pipelineDepth; ++i) {
            threadSum += (localX.f[i * fPerF4 + 0] + localX.f[i * fPerF4 + 1])
                + (localX.f[i * fPerF4 + 2] + localX.f[i * fPerF4 + 3]);
		}

        index += stride;
    }

    int remainderIndex = (vectorizedIterationCount * floatsPerChunk) + blockIdx.x * blockDim.x + threadIdx.x;
	while (remainderIndex < size) {
		threadSum += data[remainderIndex];
		remainderIndex += blockDim.x * gridDim.x;
	}

    T warpSum = warpReduction(threadSum);

    //first thread of each warp writes the warp's sum to shared
    uint warpNumber = threadIdx.x >> 5;
    if (!(threadIdx.x & 0b11111)) {
        warpSums[warpNumber] = warpSum;
    }

    __syncthreads();

    //first warp reduces shared sums
    if (threadIdx.x < blockWarpCount) {
        T blockSum = warpReduction(warpSums[threadIdx.x], blockReductionThreadMask);
        //first thread writes result to global
        if (!threadIdx.x) {
            atomicAdd(out, (T2)blockSum);
        }
    }
}




torch::Tensor kernelCaller(torch::Tensor x, torch::Tensor result) {
    //at::TensorOptions opts = at::TensorOptions().dtype(at::kDouble).device(at::kCUDA);
	//auto output = torch::zeros(1, opts); //at::TensorOptions = {}
    constexpr int threads = 512;
    constexpr int blocks = 132 * 3; //H100 has 132 SMs, 2048 max threads per sm. a100 is 108
    //const int blocks = max((x.numel() + threads - 1) / (threads * 1024), (long)1);
    //reduction_basic<float, double, threads><<<blocks, threads>>>(
    reduction_vectorized<float, double, threads, 2><<<blocks, threads>>>(
		reinterpret_cast<float*>(x.data_ptr<float>()),
        reinterpret_cast<double*>(result.data_ptr<double>()),
		x.numel());
    return result;
}
"""

# Here, the C++ source need only declare the function signature.
cpp_source = "torch::Tensor kernelCaller(torch::Tensor x, torch::Tensor result);"

module = torch.utils.cpp_extension.load_inline(
    name="module",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["kernelCaller"],
    verbose=True,
    extra_cuda_cflags=["-O2"],
)

result = torch.tensor(0.0, dtype=torch.float64, device="cuda:0")


def custom_kernel(data: input_t) -> output_t:
    # result = module.kernelCaller(data)
    data, _ = data
    result.fill_(0.0)
    module.kernelCaller(data, result)
    return result