submission 780532

#!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 = 128;
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;
}

__inline__ __device__ float sum_float4(float4 v) {
  return (v.x + v.y) + (v.z + v.w);
}

__global__ __launch_bounds__(kThreadsPerBlock, 4) 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];
  const float4* x4 = reinterpret_cast<const float4*>(x);
  const int vec_index = tile_idx * (kBlockItems / 4) + tid;

  float acc0 = 0.0f;
  float acc1 = 0.0f;
  float acc2 = 0.0f;
  float acc3 = 0.0f;

#pragma unroll
  for (int i = 0; i < kItemsPerThread / 4; i += 4) {
    float4 v0 = __ldcs(&x4[vec_index + i * kThreadsPerBlock]);
    float4 v1 = __ldcs(&x4[vec_index + (i + 1) * kThreadsPerBlock]);
    float4 v2 = __ldcs(&x4[vec_index + (i + 2) * kThreadsPerBlock]);
    float4 v3 = __ldcs(&x4[vec_index + (i + 3) * kThreadsPerBlock]);
    acc0 += sum_float4(v0);
    acc1 += sum_float4(v1);
    acc2 += sum_float4(v2);
    acc3 += sum_float4(v3);
  }

  float acc = (acc0 + acc1) + (acc2 + acc3);
  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 acc0 = 0.0f;
  float acc1 = 0.0f;
  float acc2 = 0.0f;
  float acc3 = 0.0f;

  for (int64_t index = tail_start + tid; index < n_elements; index += 4 * kThreadsPerBlock) {
    acc0 += x[index];
    if (index + kThreadsPerBlock < n_elements) {
      acc1 += x[index + kThreadsPerBlock];
    }
    if (index + 2 * kThreadsPerBlock < n_elements) {
      acc2 += x[index + 2 * kThreadsPerBlock];
    }
    if (index + 3 * kThreadsPerBlock < n_elements) {
      acc3 += x[index + 3 * kThreadsPerBlock];
    }
  }

  float acc = (acc0 + acc1) + (acc2 + acc3);
  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_v6_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]