submission 780525

#!POPCORN leaderboard vectorsum_v2

import hashlib
import os
import subprocess
import sys
from pathlib import Path

from task import input_t, output_t

from torch.utils.cpp_extension import CUDA_HOME, load_inline

subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "ninja"])

_extension_module = None

CPP_SRC = r"""
#include <torch/extension.h>

void run_vectorsum_cuda(torch::Tensor x, torch::Tensor out);

void run_vectorsum(torch::Tensor x, torch::Tensor out) {
  run_vectorsum_cuda(x, out);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("run_vectorsum", &run_vectorsum, "Vector sum CUDA launcher");
}
"""

CUDA_SRC = r"""
#include <cuda_runtime.h>
#include <torch/extension.h>

namespace {

constexpr int kThreadsPerBlock = 256;
constexpr int kWarpsPerBlock = kThreadsPerBlock / 32;
constexpr int kItemsPerThread = 32;
constexpr int kBlockItems = kThreadsPerBlock * kItemsPerThread;

__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;
}

__global__ __launch_bounds__(kThreadsPerBlock, 8) void reduce_tiles_aligned_kernel(
    const float* __restrict__ x,
    float* __restrict__ out) {
  const int tid = threadIdx.x;
  const int lane = tid & 31;
  const int warp = tid >> 5;
  const int tile_idx = blockIdx.x;

  __shared__ float warp_sums[kWarpsPerBlock];
  float acc = 0.0f;
  const float4* x4 = reinterpret_cast<const float4*>(x);

  const int vec_index = tile_idx * (kBlockItems / 4) + tid;
  float4 v0 = __ldcs(&x4[vec_index]);
  float4 v1 = __ldcs(&x4[vec_index + kThreadsPerBlock]);
  acc += v0.x + v0.y + v0.z + v0.w;
  acc += v1.x + v1.y + v1.z + v1.w;
  v0 = __ldcs(&x4[vec_index + 2 * kThreadsPerBlock]);
  v1 = __ldcs(&x4[vec_index + 3 * kThreadsPerBlock]);
  acc += v0.x + v0.y + v0.z + v0.w;
  acc += v1.x + v1.y + v1.z + v1.w;
  v0 = __ldcs(&x4[vec_index + 4 * kThreadsPerBlock]);
  v1 = __ldcs(&x4[vec_index + 5 * kThreadsPerBlock]);
  acc += v0.x + v0.y + v0.z + v0.w;
  acc += v1.x + v1.y + v1.z + v1.w;
  v0 = __ldcs(&x4[vec_index + 6 * kThreadsPerBlock]);
  v1 = __ldcs(&x4[vec_index + 7 * kThreadsPerBlock]);
  acc += v0.x + v0.y + v0.z + v0.w;
  acc += v1.x + v1.y + v1.z + v1.w;

  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) {
      atomicAdd(out, block_sum);
    }
  }
}

__global__ void reduce_tail_kernel(
    const float* __restrict__ x,
    float* __restrict__ out,
    int64_t tail_start,
    int64_t n_elements) {
  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 = tail_start + tid; index < n_elements; index += kThreadsPerBlock) {
    acc += x[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) {
      atomicAdd(out, block_sum);
    }
  }
}

}  // namespace

void run_vectorsum_cuda(torch::Tensor x, torch::Tensor out) {
  (void)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;

  if (n_full_tiles > 0) {
    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>(),
        out.data_ptr<float>(),
        full_tile_elements,
        n_elements);
  }
}
"""


def _module_name() -> str:
    source_hash = hashlib.sha256((CPP_SRC + CUDA_SRC).encode("utf-8")).hexdigest()[:16]
    return f"vectorsum_v5_ext_{source_hash}"


def _load_extension():
    global _extension_module

    if _extension_module is not None:
        return _extension_module

    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,
    )
    return _extension_module


ext = _load_extension()

def custom_kernel(data: input_t) -> output_t:
    vector, output = data
    ext.run_vectorsum(vector, output)
    return output[0]