submission 780282

  ⋯ 12 unchanged lines
   
  subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "ninja"])
   
- THREADS_PER_BLOCK = 256
- WARPS_PER_BLOCK = THREADS_PER_BLOCK // 32
- ITEMS_PER_THREAD = 16
- BLOCK_ITEMS = THREADS_PER_BLOCK * ITEMS_PER_THREAD
  _extension_module = None
- _partial_buffer_cache: dict[int, torch.Tensor] = {}
   
  CPP_SRC = r"""
  #include <torch/extension.h>
   
- void run_vectorsum_cuda(torch::Tensor x, torch::Tensor partials, torch::Tensor out);
+ void run_vectorsum_cuda(torch::Tensor x, torch::Tensor out);
   
- void run_vectorsum(torch::Tensor x, torch::Tensor partials, torch::Tensor out) {
+ void run_vectorsum(torch::Tensor x, torch::Tensor out) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
-   TORCH_CHECK(partials.is_cuda(), "partials must be a CUDA tensor");
    TORCH_CHECK(out.is_cuda(), "out must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "x must be float32");
-   TORCH_CHECK(partials.scalar_type() == torch::kFloat32, "partials must be float32");
    TORCH_CHECK(out.scalar_type() == torch::kFloat32, "out must be float32");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
-   TORCH_CHECK(partials.is_contiguous(), "partials must be contiguous");
    TORCH_CHECK(out.is_contiguous(), "out must be contiguous");
    TORCH_CHECK(x.dim() == 1, "x must be 1D");
-   TORCH_CHECK(partials.dim() == 1, "partials must be 1D");
    TORCH_CHECK(out.numel() == 1, "out must have one element");
-   TORCH_CHECK(x.device() == partials.device(), "x and partials must be on same device");
    TORCH_CHECK(x.device() == out.device(), "x and out must be on same device");
-   run_vectorsum_cuda(x, partials, out);
+   run_vectorsum_cuda(x, out);
  }
   
  PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  ⋯ 3 unchanged lines
   
  CUDA_SRC = r"""
  #include <c10/cuda/CUDAGuard.h>
+ #include <c10/cuda/CUDAException.h>
  #include <cuda_runtime.h>
  #include <torch/extension.h>
   
  ⋯ 1 unchanged lines
   
  constexpr int kThreadsPerBlock = 256;
  constexpr int kWarpsPerBlock = kThreadsPerBlock / 32;
- constexpr int kItemsPerThread = 16;
+ constexpr int kItemsPerThread = 32;
  constexpr int kBlockItems = kThreadsPerBlock * kItemsPerThread;
   
  __inline__ __device__ float warp_reduce_sum(float val) {
  ⋯ 3 unchanged lines
    return val;
  }
   
- __global__ __launch_bounds__(kThreadsPerBlock, 8) void reduce_tiles_fast_aligned_kernel(
+ __global__ __launch_bounds__(kThreadsPerBlock, 8) void reduce_tiles_aligned_kernel(
      const float* __restrict__ x,
-     float* __restrict__ partials,
-     int n_full_tiles) {
+     float* __restrict__ out) {
    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
-   const int grid_stride_tiles = gridDim.x;
+   const int tile_idx = blockIdx.x;
   
    __shared__ float warp_sums[kWarpsPerBlock];
-   float acc0 = 0.0f;
-   float acc1 = 0.0f;
-   float acc2 = 0.0f;
-   float acc3 = 0.0f;
-  
+   float acc = 0.0f;
    const float4* x4 = reinterpret_cast<const float4*>(x);
   
-   for (int tile_idx = blockIdx.x; tile_idx < n_full_tiles; tile_idx += grid_stride_tiles) {
-     const int vec_index = tile_idx * (kBlockItems / 4) + tid;
-     float4 v0 = x4[vec_index];
-     float4 v1 = x4[vec_index + kThreadsPerBlock];
-     float4 v2 = x4[vec_index + 2 * kThreadsPerBlock];
-     float4 v3 = x4[vec_index + 3 * kThreadsPerBlock];
-     acc0 += v0.x + v0.y + v0.z + v0.w;
-     acc1 += v1.x + v1.y + v1.z + v1.w;
-     acc2 += v2.x + v2.y + v2.z + v2.w;
-     acc3 += v3.x + v3.y + v3.z + v3.w;
-   }
+   const int vec_index = tile_idx * (kBlockItems / 4) + tid;
+   float4 v = x4[vec_index];
+   acc += v.x + v.y + v.z + v.w;
+   v = x4[vec_index + kThreadsPerBlock];
+   acc += v.x + v.y + v.z + v.w;
+   v = x4[vec_index + 2 * kThreadsPerBlock];
+   acc += v.x + v.y + v.z + v.w;
+   v = x4[vec_index + 3 * kThreadsPerBlock];
+   acc += v.x + v.y + v.z + v.w;
+   v = x4[vec_index + 4 * kThreadsPerBlock];
+   acc += v.x + v.y + v.z + v.w;
+   v = x4[vec_index + 5 * kThreadsPerBlock];
+   acc += v.x + v.y + v.z + v.w;
+   v = x4[vec_index + 6 * kThreadsPerBlock];
+   acc += v.x + v.y + v.z + v.w;
+   v = x4[vec_index + 7 * kThreadsPerBlock];
+   acc += v.x + v.y + v.z + v.w;
   
-   float acc = (acc0 + acc1) + (acc2 + acc3);
    acc = warp_reduce_sum(acc);
    if (lane == 0) {
      warp_sums[warp] = acc;
  ⋯ 4 unchanged lines
      float block_sum = (lane < kWarpsPerBlock) ? warp_sums[lane] : 0.0f;
      block_sum = warp_reduce_sum(block_sum);
      if (lane == 0) {
-       partials[blockIdx.x] = block_sum;
+       atomicAdd(out, block_sum);
      }
    }
  }
   
- __global__ void reduce_tiles_scalar_kernel(
-     const float* __restrict__ x,
-     float* __restrict__ partials,
-     int n_full_tiles) {
-   const int tid = threadIdx.x;
-   const int lane = tid & 31;
-   const int warp = tid >> 5;
-   const int grid_stride_tiles = gridDim.x;
-  
-   __shared__ float warp_sums[kWarpsPerBlock];
-   float acc = 0.0f;
-  
-   for (int tile_idx = blockIdx.x; tile_idx < n_full_tiles; tile_idx += grid_stride_tiles) {
-     const int tile_base = tile_idx * kBlockItems + tid;
-     acc += x[tile_base];
-     acc += x[tile_base + kThreadsPerBlock];
-     acc += x[tile_base + 2 * kThreadsPerBlock];
-     acc += x[tile_base + 3 * kThreadsPerBlock];
-     acc += x[tile_base + 4 * kThreadsPerBlock];
-     acc += x[tile_base + 5 * kThreadsPerBlock];
-     acc += x[tile_base + 6 * kThreadsPerBlock];
-     acc += x[tile_base + 7 * kThreadsPerBlock];
-     acc += x[tile_base + 8 * kThreadsPerBlock];
-     acc += x[tile_base + 9 * kThreadsPerBlock];
-     acc += x[tile_base + 10 * kThreadsPerBlock];
-     acc += x[tile_base + 11 * kThreadsPerBlock];
-     acc += x[tile_base + 12 * kThreadsPerBlock];
-     acc += x[tile_base + 13 * kThreadsPerBlock];
-     acc += x[tile_base + 14 * kThreadsPerBlock];
-     acc += x[tile_base + 15 * kThreadsPerBlock];
-   }
-  
-   acc = warp_reduce_sum(acc);
-   if (lane == 0) {
-     warp_sums[warp] = acc;
-   }
-   __syncthreads();
-  
-   if (warp == 0) {
-     float block_sum = (lane < kWarpsPerBlock) ? warp_sums[lane] : 0.0f;
-     block_sum = warp_reduce_sum(block_sum);
-     if (lane == 0) {
-       partials[blockIdx.x] = block_sum;
-     }
-   }
- }
-  
  __global__ void reduce_tail_kernel(
      const float* __restrict__ x,
-     float* __restrict__ partials,
+     float* __restrict__ out,
      int64_t tail_start,
-     int64_t n_elements,
-     int tail_partial_index) {
+     int64_t n_elements) {
    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
  ⋯ 15 unchanged lines
      float block_sum = (lane < kWarpsPerBlock) ? warp_sums[lane] : 0.0f;
      block_sum = warp_reduce_sum(block_sum);
      if (lane == 0) {
-       partials[tail_partial_index] = block_sum;
+       atomicAdd(out, block_sum);
      }
    }
  }
   
- __global__ void final_reduce_kernel(
-     const float* __restrict__ partials,
-     float* __restrict__ out,
-     int64_t n_partials) {
-   const int tid = threadIdx.x;
-   const int lane = tid & 31;
-   const int warp = tid >> 5;
-  
-   __shared__ float warp_sums[kWarpsPerBlock];
-   float acc = 0.0f;
-  
-   for (int64_t index = tid; index < n_partials; index += kThreadsPerBlock) {
-     acc += partials[index];
-   }
-  
-   acc = warp_reduce_sum(acc);
-   if (lane == 0) {
-     warp_sums[warp] = acc;
-   }
-   __syncthreads();
-  
-   if (warp == 0) {
-     float block_sum = (lane < kWarpsPerBlock) ? warp_sums[lane] : 0.0f;
-     block_sum = warp_reduce_sum(block_sum);
-     if (lane == 0) {
-       out[0] = block_sum;
-     }
-   }
- }
-  
  }  // namespace
   
- void run_vectorsum_cuda(torch::Tensor x, torch::Tensor partials, torch::Tensor out) {
+ void run_vectorsum_cuda(torch::Tensor x, torch::Tensor out) {
    const c10::cuda::CUDAGuard device_guard(x.device());
+   C10_CUDA_CHECK(cudaMemsetAsync(out.data_ptr<float>(), 0, sizeof(float)));
   
    const int64_t n_elements = x.numel();
    const int64_t n_full_tiles = n_elements / kBlockItems;
    const int64_t full_tile_elements = n_full_tiles * kBlockItems;
    const bool has_tail = full_tile_elements < n_elements;
-   const int64_t n_partials = partials.numel();
-   const bool aligned = ((reinterpret_cast<uintptr_t>(x.data_ptr<float>()) & 0xf) == 0);
   
-   int partial_index = 0;
    if (n_full_tiles > 0) {
-     const int full_tile_blocks = has_tail ? static_cast<int>(n_partials - 1) : static_cast<int>(n_partials);
-     const int n_full_tiles_int = static_cast<int>(n_full_tiles);
-     if (aligned) {
-       reduce_tiles_fast_aligned_kernel<<<full_tile_blocks, kThreadsPerBlock>>>(
-           x.data_ptr<float>(),
-           partials.data_ptr<float>(),
-           n_full_tiles_int);
-     } else {
-       reduce_tiles_scalar_kernel<<<full_tile_blocks, kThreadsPerBlock>>>(
-           x.data_ptr<float>(),
-           partials.data_ptr<float>(),
-           n_full_tiles_int);
-     }
-     partial_index = full_tile_blocks;
+     const int full_tile_blocks = static_cast<int>(n_full_tiles);
+     reduce_tiles_aligned_kernel<<<full_tile_blocks, kThreadsPerBlock>>>(
+         x.data_ptr<float>(),
+         out.data_ptr<float>());
    }
    if (has_tail) {
      reduce_tail_kernel<<<1, kThreadsPerBlock>>>(
          x.data_ptr<float>(),
-         partials.data_ptr<float>(),
+         out.data_ptr<float>(),
          full_tile_elements,
-         n_elements,
-         partial_index);
+         n_elements);
    }
-   final_reduce_kernel<<<1, kThreadsPerBlock>>>(
-       partials.data_ptr<float>(),
-       out.data_ptr<float>(),
-       n_partials);
  }
  """
   
   
- def _ceil_div(x: int, y: int) -> int:
-     return (x + y - 1) // y
-  
-  
- def _rounded_capacity(size: int) -> int:
-     capacity = 1
-     while capacity < size:
-         capacity *= 2
-     return capacity
-  
-  
  def _module_name() -> str:
      source_hash = hashlib.sha256((CPP_SRC + CUDA_SRC).encode("utf-8")).hexdigest()[:16]
-     return f"vectorsum_v3_ext_{source_hash}"
+     return f"vectorsum_v4_ext_{source_hash}"
   
   
  def _load_extension():
  ⋯ 21 unchanged lines
      return _extension_module
   
   
- def _get_partial_buffer(device: torch.device, num_blocks: int) -> torch.Tensor:
-     device_index = device.index
-     if device_index is None:
-         device_index = torch.cuda.current_device()
-  
-     partial_buffer = _partial_buffer_cache.get(device_index)
-     if partial_buffer is None or partial_buffer.numel() < num_blocks:
-         partial_buffer = torch.empty(
-             _rounded_capacity(num_blocks),
-             device=device,
-             dtype=torch.float32,
-         )
-         _partial_buffer_cache[device_index] = partial_buffer
-     return partial_buffer[:num_blocks]
-  
-  
- def _num_blocks(n_elements: int) -> int:
-     return max(1, _ceil_div(n_elements, BLOCK_ITEMS))
-  
-  
  def custom_kernel(data: input_t) -> output_t:
      vector, output = data
      n_elements = vector.numel()
  ⋯ 3 unchanged lines
          return output[0]
   
      ext = _load_extension()
-     n_full_tiles = n_elements // BLOCK_ITEMS
-     full_tile_elements = n_full_tiles * BLOCK_ITEMS
-     has_tail = full_tile_elements < n_elements
-     num_full_blocks = _num_blocks(full_tile_elements) if n_full_tiles > 0 else 0
-     num_partials = num_full_blocks + (1 if has_tail else 0)
-     partials = _get_partial_buffer(vector.device, num_partials)
-     ext.run_vectorsum(vector, partials, output)
+     ext.run_vectorsum(vector, output)
      return output[0]