submission 68130

  #!POPCORN leaderboard vectorsum_v2
+ #!POPCORN gpus 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
   
- import functools
+ cuda_source = """
+ #include <pybind11/pybind11.h>
+ //#include <pybind11/tuple.h>
+ #include <torch/extension.h>
   
- try:
-     import cuda.parallel.experimental.algorithms as algorithms
- except:
-     import os
-     import subprocess
+ 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;
+ }
   
-     gpu = "a100"
+ template <typename T, typename T2, uint blockSize>
+ __global__ void reduction_basic(T* data, T2* out, size_t size) {
+     uint index = blockIdx.x * blockDim.x + threadIdx.x;
+     const uint stride = gridDim.x * blockDim.x;
   
-     if not os.path.exists("cccl"):
-         subprocess.check_call(
-             ["git", "clone", "https://github.com/NaderAlAwar/cccl.git"]
-         )
-     subprocess.check_call(["git", "checkout", "gpu-mode-submissions-a100"], cwd="cccl/")
-     subprocess.check_call(
-         ["git", "pull", "origin", "gpu-mode-submissions-a100"], cwd="cccl/"
-     )
-     if gpu != "l4":
-         subprocess.check_call(
-             ["git", "checkout", "6092a0bead3a297666a2aba37672f67a8a1a569b"],
-             cwd="cccl/python/cuda_parallel",
-         )
-     else:
-         subprocess.check_call(
-             ["git", "checkout", "bc36bca5c76e3e9c6bf78d6906127e8549cb551a"],
-             cwd="cccl/python/cuda_parallel",
-         )
+     constexpr uint blockWarpCount = (blockSize + 31) / 32;
+     constexpr uint blockReductionThreadMask = (1 << blockWarpCount) - 1;
+     __shared__ T warpSums[blockWarpCount];
   
-     env = os.environ.copy()
-     env["CC"] = "gcc"
-     env["CXX"] = "g++"
-     env["CMAKE_ARGS"] = "-DCMAKE_CXX_STANDARD=20"
+     T threadSum {};
+     while (index < size) {
+         threadSum += data[index];
+         index += stride;
+     }
   
-     subprocess.check_call(
-         ["pip", "install", "../cuda_cccl"], cwd="cccl/python/cuda_parallel", env=env
-     )
-     subprocess.check_call(
-         ["pip", "install", ".[test]", "-v"], cwd="cccl/python/cuda_parallel", env=env
-     )
-     subprocess.check_call(
-         ["pip", "install", "cupy-cuda12x"], cwd="cccl/python/cuda_parallel", env=env
-     )
+     T warpSum = warpReduction(threadSum);
   
- import cupy as cp
- import numpy as np
- import cuda.parallel.experimental.algorithms as algorithms
- import functools
+     //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();
   
- def add_op(a, b):
-     return a + b
+     //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 = 256;
+     const int blocks = (x.numel() + threads - 1) / threads;
+     reduction_basic<float, double, 256><<<blocks, threads>>>(
+ 		reinterpret_cast<float*>(x.data_ptr<float>()),
+         reinterpret_cast<double*>(result.data_ptr<double>()),
+ 		x.numel());
+     return result;
+ }
+ """
   
- import torch
+ # Here, the C++ source need only declare the function signature.
+ cpp_source = "torch::Tensor kernelCaller(torch::Tensor x, torch::Tensor result);"
   
- from task import input_t, output_t
+ 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")
   
- @functools.cache
- def initialize(num_items):
-     d_in = torch.tensor(num_items, dtype=torch.float32).cuda()
-     d_out = torch.tensor(0, dtype=torch.float32).cuda()
-     h_init = np.array([0], dtype="float32")
-     reducer = algorithms.nondeterministic_reduce_into(d_in, d_out, add_op, h_init)
-     temp_storage_size = reducer(None, d_in, d_out, num_items, h_init)
   
-     d_temp_storage = cp.empty(temp_storage_size, dtype=np.uint8)
-     # reducer.initialize_fast(d_temp_storage, d_in, d_out, num_items, h_init)
-  
-     return d_temp_storage, d_out, h_init, reducer
-  
-  
  def custom_kernel(data: input_t) -> output_t:
+     # result = module.kernelCaller(data)
      data, _ = data
-     num_items = data.shape[0]
-     d_temp_storage, d_out, h_init, reducer = initialize(num_items)
-     reducer(d_temp_storage, data, d_out, num_items, h_init)
-     # _, d_out, _, reducer = initialize(num_items)
-     # reducer.call_fast(data)
-  
-     return d_out
+     result.fill_(0.0)
+     module.kernelCaller(data, result)
+     return result