submission 755314

  #!POPCORN leaderboard vectorsum_v2
   
- import os
- os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
-  
  import torch
- from torch.utils.cpp_extension import load_inline
+ import triton
+ import triton.language as tl
  from task import input_t, output_t
   
- cuda_source = r"""
- #include <cuda_runtime.h>
- #include <torch/extension.h>
   
- __device__ __forceinline__ double warp_reduce(double val) {
-     val += __shfl_down_sync(0xffffffff, val, 16);
-     val += __shfl_down_sync(0xffffffff, val, 8);
-     val += __shfl_down_sync(0xffffffff, val, 4);
-     val += __shfl_down_sync(0xffffffff, val, 2);
-     val += __shfl_down_sync(0xffffffff, val, 1);
-     return val;
- }
+ @triton.jit
+ def sum_kernel_single(
+ 	x_ptr,
+ 	partial_ptr,
+ 	counter_ptr,
+ 	output_ptr,
+ 	n_elements,
+ 	n_blocks,
+ 	BLOCK_SIZE: tl.constexpr,
+ 	FINAL_BLOCK: tl.constexpr,
+ ):
+ 	pid = tl.program_id(0)
+ 	block_start = pid * BLOCK_SIZE
+ 	offsets = block_start + tl.arange(0, BLOCK_SIZE)
+ 	mask = offsets < n_elements
+ 	x = tl.load(x_ptr + offsets, mask=mask, other=0.0, eviction_policy="evict_first")
+ 	block_sum = tl.sum(x, axis=0)
+ 	tl.store(partial_ptr + pid, block_sum)
+ 	tl.debug_barrier()
+ 	old = tl.atomic_add(counter_ptr, 1)
+ 	is_last = old == (n_blocks - 1)
+ 	if is_last:
+ 		offsets2 = tl.arange(0, FINAL_BLOCK)
+ 		mask2 = offsets2 < n_blocks
+ 		vals = tl.load(partial_ptr + offsets2, mask=mask2, other=0.0)
+ 		total = tl.sum(vals, axis=0)
+ 		tl.store(output_ptr, total)
+ 		tl.store(counter_ptr, 0)
   
- __device__ __forceinline__ float4 load_cs(const float4* ptr) {
-     float4 ret;
-     asm volatile("ld.global.cs.v4.f32 {%0, %1, %2, %3}, [%4];"
-         : "=f"(ret.x), "=f"(ret.y), "=f"(ret.z), "=f"(ret.w)
-         : "l"(ptr));
-     return ret;
- }
   
- __global__ void __launch_bounds__(256)
- vectorsum_kernel(
-     const float* __restrict__ input,
-     double* __restrict__ acc,
-     const int n
- ) {
-     double s0 = 0.0, s1 = 0.0, s2 = 0.0, s3 = 0.0;
-     const int tid = threadIdx.x + blockIdx.x * 256;
-     const int gs = 256 * gridDim.x;
-     const int n4 = n >> 2;
-     const float4* in4 = reinterpret_cast<const float4*>(input);
+ _COUNTER = torch.zeros(1, device="cuda", dtype=torch.int32)
+ _PARTS = torch.empty(12800, device="cuda", dtype=torch.float32)
   
-     int i = tid;
-     for (; i + gs * 3 < n4; i += gs * 4) {
-         float4 v0 = load_cs(&in4[i]);
-         float4 v1 = load_cs(&in4[i + gs]);
-         float4 v2 = load_cs(&in4[i + gs * 2]);
-         float4 v3 = load_cs(&in4[i + gs * 3]);
-         s0 += (double)v0.x + (double)v0.y + (double)v0.z + (double)v0.w;
-         s1 += (double)v1.x + (double)v1.y + (double)v1.z + (double)v1.w;
-         s2 += (double)v2.x + (double)v2.y + (double)v2.z + (double)v2.w;
-         s3 += (double)v3.x + (double)v3.y + (double)v3.z + (double)v3.w;
-     }
-     for (; i < n4; i += gs) {
-         float4 v = load_cs(&in4[i]);
-         s0 += (double)v.x + (double)v.y + (double)v.z + (double)v.w;
-     }
-     for (int j = (n4 << 2) + tid; j < n; j += gs)
-         s0 += (double)__ldg(&input[j]);
-  
-     double ts = (s0 + s1) + (s2 + s3);
-     ts = warp_reduce(ts);
-  
-     __shared__ double ws[8];
-     const int lane = threadIdx.x & 31;
-     const int wid = threadIdx.x >> 5;
-     if (lane == 0) ws[wid] = ts;
-     __syncthreads();
-     if (wid == 0) {
-         double v = (lane < 8) ? ws[lane] : 0.0;
-         v = warp_reduce(v);
-         if (lane == 0) atomicAdd(acc, v);
-     }
+ # (BLOCK_SIZE, n_blocks, FINAL_BLOCK, num_warps)
+ _TABLE = {}
+ _RAW = {
+ 	1638400: (16384, 16),
+ 	3276800: (8192, 8),
+ 	6553600: (8192, 4),
+ 	13107200: (8192, 4),
+ 	26214400: (32768, 8),
+ 	52428800: (32768, 16),
  }
   
- __global__ void finalize(double* __restrict__ acc, float* __restrict__ out) {
-     *out = (float)(*acc);
-     *acc = 0.0;
- }
+ for _n, (_bs, _nw) in _RAW.items():
+ 	_nb = (_n + _bs - 1) // _bs
+ 	_b2 = max(32, 1 << (_nb - 1).bit_length())
+ 	_TABLE[_n] = (_bs, _nb, _b2, _nw)
   
- void vectorsum(torch::Tensor input, torch::Tensor output,
-                torch::Tensor acc, int nb) {
-     vectorsum_kernel<<<nb, 256>>>(
-         input.data_ptr<float>(), acc.data_ptr<double>(), input.numel());
-     finalize<<<1, 1>>>(acc.data_ptr<double>(), output.data_ptr<float>());
- }
- """
   
- cpp_source = "void vectorsum(torch::Tensor, torch::Tensor, torch::Tensor, int);"
+ def _get_config(n):
+ 	if n in _TABLE:
+ 		return _TABLE[n]
+ 	# Fallback for unseen sizes
+ 	bs = 8192
+ 	nw = 8
+ 	nb = (n + bs - 1) // bs
+ 	b2 = max(32, 1 << (nb - 1).bit_length())
+ 	return (bs, nb, b2, nw)
   
- module = load_inline(
-     name="vectorsum_cuda",
-     cpp_sources=[cpp_source],
-     cuda_sources=[cuda_source],
-     functions=["vectorsum"],
-     extra_cuda_cflags=["-O3", "--use_fast_math"],
-     verbose=False,
- )
   
- _acc = torch.zeros(1, dtype=torch.float64, device="cuda")
-  
-  
- def _tune():
-     import sys
-     sys.path.insert(0, '.')
-     from reference import generate_input
-     size = 52428800
-     best_t, best_nb = 1e9, 5920
-     for nb in [2960, 4096, 5920]:
-         data = generate_input(size=size, seed=12345)
-         for _ in range(3):
-             module.vectorsum(data[0], data[1], _acc, nb)
-             torch.cuda.synchronize()
-         ts = []
-         for _ in range(15):
-             data = generate_input(size=size, seed=12345)
-             torch.cuda.synchronize()
-             s, e = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
-             s.record(); module.vectorsum(data[0], data[1], _acc, nb); e.record()
-             torch.cuda.synchronize(); ts.append(s.elapsed_time(e))
-         a = sum(ts)/len(ts)
-         if a < best_t: best_t, best_nb = a, nb
-     return best_nb
-  
- _nb = _tune()
-  
-  
  def custom_kernel(data: input_t) -> output_t:
-     inp, out = data
-     n = inp.numel()
-     nb = max(1, min(_nb, n // 64))
-     module.vectorsum(inp, out, _acc, nb)
-     return out[0]
+ 	data, output = data
+ 	n = data.numel()
+ 	bs, nb, b2, nw = _get_config(n)
+ 	sum_kernel_single[(nb,)](
+ 		data, _PARTS, _COUNTER, output, n, nb,
+ 		BLOCK_SIZE=bs, FINAL_BLOCK=b2, num_warps=nw,
+ 	)
+ 	return output[0]