submission 780050

  #!POPCORN leaderboard vectorsum_v2
   
-  
+ import hashlib
+ import os
  import subprocess
  import sys
+ from pathlib import Path
   
- for pkg in ["apache-tvm-ffi", "torch-c-dlpack-ext"]:
-     subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "--upgrade", pkg])
-  
- from collections.abc import Callable
-  
  from task import input_t, output_t
   
- import cutlass
- import cutlass.cute as cute
- import cutlass.utils as utils
  import torch
+ from torch.utils.cpp_extension import CUDA_HOME, load_inline
   
+ subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "ninja"])
+  
  THREADS_PER_BLOCK = 256
  WARPS_PER_BLOCK = THREADS_PER_BLOCK // 32
- ITEMS_PER_THREAD = 8
+ ITEMS_PER_THREAD = 16
  BLOCK_ITEMS = THREADS_PER_BLOCK * ITEMS_PER_THREAD
- MAX_BLOCKS_PER_SM = 8
-  
- _compiled_reduce_tiles = None
- _compiled_final_reduce = None
+ _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);
   
- @cute.kernel
- def reduce_tiles_kernel(vector: cute.Tensor, partials: cute.Tensor):
-     tid_x, _, _ = cute.arch.thread_idx()
-     cta_x, _, _ = cute.arch.block_idx()
-     grid_x, _, _ = cute.arch.grid_dim()
-     lane = cute.arch.lane_idx()
-     warp = cute.arch.warp_idx()
+ void run_vectorsum(torch::Tensor x, torch::Tensor partials, 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);
+ }
   
-     n_elements = vector.shape[0]
-     grid_stride = grid_x * BLOCK_ITEMS
-     base_index = cta_x * BLOCK_ITEMS + tid_x
+ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+   m.def("run_vectorsum", &run_vectorsum, "Vector sum CUDA launcher");
+ }
+ """
   
-     smem = utils.SmemAllocator()
-     warp_sums = smem.allocate_tensor(cutlass.Float32, WARPS_PER_BLOCK)
+ CUDA_SRC = r"""
+ #include <c10/cuda/CUDAGuard.h>
+ #include <cuda_runtime.h>
+ #include <torch/extension.h>
   
-     acc = cutlass.Float32(0.0)
-     index = base_index
-     while index < n_elements:
-         idx0 = index
-         idx1 = index + THREADS_PER_BLOCK
-         idx2 = index + 2 * THREADS_PER_BLOCK
-         idx3 = index + 3 * THREADS_PER_BLOCK
-         idx4 = index + 4 * THREADS_PER_BLOCK
-         idx5 = index + 5 * THREADS_PER_BLOCK
-         idx6 = index + 6 * THREADS_PER_BLOCK
-         idx7 = index + 7 * THREADS_PER_BLOCK
+ namespace {
   
-         if idx0 < n_elements:
-             acc = acc + vector[idx0]
-         if idx1 < n_elements:
-             acc = acc + vector[idx1]
-         if idx2 < n_elements:
-             acc = acc + vector[idx2]
-         if idx3 < n_elements:
-             acc = acc + vector[idx3]
-         if idx4 < n_elements:
-             acc = acc + vector[idx4]
-         if idx5 < n_elements:
-             acc = acc + vector[idx5]
-         if idx6 < n_elements:
-             acc = acc + vector[idx6]
-         if idx7 < n_elements:
-             acc = acc + vector[idx7]
+ constexpr int kThreadsPerBlock = 256;
+ constexpr int kWarpsPerBlock = kThreadsPerBlock / 32;
+ constexpr int kItemsPerThread = 16;
+ constexpr int kBlockItems = kThreadsPerBlock * kItemsPerThread;
   
-         index = index + grid_stride
+ __inline__ __device__ float warp_reduce_sum(float val) {
+   for (int offset = 16; offset > 0; offset /= 2) {
+     val += __shfl_down_sync(0xffffffff, val, offset);
+   }
+   return val;
+ }
   
-     acc = cute.arch.warp_reduction_sum(acc)
-     if lane == 0:
-         warp_sums[warp] = acc
+ __global__ __launch_bounds__(kThreadsPerBlock, 8) void reduce_tiles_fast_aligned_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;
   
-     cute.arch.sync_threads()
+   __shared__ float warp_sums[kWarpsPerBlock];
+   float acc0 = 0.0f;
+   float acc1 = 0.0f;
+   float acc2 = 0.0f;
+   float acc3 = 0.0f;
   
-     if warp == 0:
-         block_sum = warp_sums[lane] if lane < WARPS_PER_BLOCK else cutlass.Float32(0.0)
-         block_sum = cute.arch.warp_reduction_sum(block_sum)
-         if lane == 0:
-             partials[cta_x] = block_sum
+   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;
+   }
   
- @cute.kernel
- def final_reduce_kernel(partials: cute.Tensor, output: cute.Tensor):
-     tid_x, _, _ = cute.arch.thread_idx()
-     lane = cute.arch.lane_idx()
-     warp = cute.arch.warp_idx()
+   float acc = (acc0 + acc1) + (acc2 + acc3);
+   acc = warp_reduce_sum(acc);
+   if (lane == 0) {
+     warp_sums[warp] = acc;
+   }
+   __syncthreads();
   
-     n_partials = partials.shape[0]
+   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;
+     }
+   }
+ }
   
-     smem = utils.SmemAllocator()
-     warp_sums = smem.allocate_tensor(cutlass.Float32, WARPS_PER_BLOCK)
+ __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;
   
-     acc = cutlass.Float32(0.0)
-     index = tid_x
-     while index < n_partials:
-         acc = acc + partials[index]
-         index = index + THREADS_PER_BLOCK
+   __shared__ float warp_sums[kWarpsPerBlock];
+   float acc = 0.0f;
   
-     acc = cute.arch.warp_reduction_sum(acc)
-     if lane == 0:
-         warp_sums[warp] = acc
+   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];
+   }
   
-     cute.arch.sync_threads()
+   acc = warp_reduce_sum(acc);
+   if (lane == 0) {
+     warp_sums[warp] = acc;
+   }
+   __syncthreads();
   
-     if warp == 0:
-         block_sum = warp_sums[lane] if lane < WARPS_PER_BLOCK else cutlass.Float32(0.0)
-         block_sum = cute.arch.warp_reduction_sum(block_sum)
-         if lane == 0:
-             output[0] = block_sum
+   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,
+     int64_t tail_start,
+     int64_t n_elements,
+     int tail_partial_index) {
+   const int tid = threadIdx.x;
+   const int lane = tid & 31;
+   const int warp = tid >> 5;
   
- @cute.jit
- def launch_reduce_tiles(vector: cute.Tensor, partials: cute.Tensor):
-     reduce_tiles_kernel(vector, partials).launch(
-         grid=(partials.shape[0], 1, 1),
-         block=[THREADS_PER_BLOCK, 1, 1],
-         cluster=(1, 1, 1),
-     )
+   __shared__ float warp_sums[kWarpsPerBlock];
+   float acc = 0.0f;
   
+   for (int64_t index = tail_start + tid; index < n_elements; index += kThreadsPerBlock) {
+     acc += x[index];
+   }
   
- @cute.jit
- def launch_final_reduce(partials: cute.Tensor, output: cute.Tensor):
-     final_reduce_kernel(partials, output).launch(
-         grid=(1, 1, 1),
-         block=[THREADS_PER_BLOCK, 1, 1],
-         cluster=(1, 1, 1),
-     )
+   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[tail_partial_index] = 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) {
+   const c10::cuda::CUDAGuard device_guard(x.device());
+  
+   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;
+   }
+   if (has_tail) {
+     reduce_tail_kernel<<<1, kThreadsPerBlock>>>(
+         x.data_ptr<float>(),
+         partials.data_ptr<float>(),
+         full_tile_elements,
+         n_elements,
+         partial_index);
+   }
+   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
   
  ⋯ 5 unchanged lines
      return capacity
   
   
- def _compile_kernels() -> tuple[Callable, Callable]:
-     global _compiled_reduce_tiles, _compiled_final_reduce
+ def _module_name() -> str:
+     source_hash = hashlib.sha256((CPP_SRC + CUDA_SRC).encode("utf-8")).hexdigest()[:16]
+     return f"vectorsum_v3_ext_{source_hash}"
   
-     if _compiled_reduce_tiles is not None and _compiled_final_reduce is not None:
-         return _compiled_reduce_tiles, _compiled_final_reduce
   
-     vector_len = cute.sym_int()
-     partial_len = cute.sym_int()
+ def _load_extension():
+     global _extension_module
   
-     fake_vector = cute.runtime.make_fake_compact_tensor(cute.Float32, (vector_len,))
-     fake_partials = cute.runtime.make_fake_compact_tensor(cute.Float32, (partial_len,))
-     fake_output = cute.runtime.make_fake_compact_tensor(cute.Float32, (1,))
+     if _extension_module is not None:
+         return _extension_module
   
-     _compiled_reduce_tiles = cute.compile(
-         launch_reduce_tiles,
-         fake_vector,
-         fake_partials,
-         options="--enable-tvm-ffi",
+     if CUDA_HOME:
+         nvcc_dir = str(Path(CUDA_HOME) / "bin")
+         path_parts = os.environ.get("PATH", "").split(os.pathsep)
+         if nvcc_dir not in path_parts:
+             os.environ["PATH"] = os.pathsep.join([nvcc_dir, *path_parts])
+  
+     _extension_module = load_inline(
+         name=_module_name(),
+         cpp_sources=CPP_SRC,
+         cuda_sources=CUDA_SRC,
+         functions=None,
+         extra_cflags=["-O3", "-std=c++17"],
+         extra_cuda_cflags=["-O3", "-std=c++17", "--use_fast_math", "-lineinfo"],
+         with_cuda=True,
+         verbose=False,
      )
-     _compiled_final_reduce = cute.compile(
-         launch_final_reduce,
-         fake_partials,
-         fake_output,
-         options="--enable-tvm-ffi",
-     )
-     return _compiled_reduce_tiles, _compiled_final_reduce
+     return _extension_module
   
   
  def _get_partial_buffer(device: torch.device, num_blocks: int) -> torch.Tensor:
  ⋯ 12 unchanged lines
      return partial_buffer[:num_blocks]
   
   
- def _num_blocks(n_elements: int, device: torch.device) -> int:
-     sm_count = torch.cuda.get_device_properties(device).multi_processor_count
-     return max(1, min(_ceil_div(n_elements, BLOCK_ITEMS), sm_count * MAX_BLOCKS_PER_SM))
+ def _num_blocks(n_elements: int) -> int:
+     return max(1, _ceil_div(n_elements, BLOCK_ITEMS))
   
   
  def custom_kernel(data: input_t) -> output_t:
  ⋯ 4 unchanged lines
          output[0] = 0.0
          return output[0]
   
-     num_blocks = _num_blocks(n_elements, vector.device)
-     partials = _get_partial_buffer(vector.device, num_blocks)
-     compiled_reduce_tiles, compiled_final_reduce = _compile_kernels()
-  
-     compiled_reduce_tiles(vector, partials)
-     compiled_final_reduce(partials, output)
-  
+     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)
      return output[0]