submission 512604

- #!POPCORN leaderboard prefixsum_v2
+ # prefixsum_v2 submission — two-pass inclusive prefix sum using Triton
+ #
+ # Algorithm:
+ #   Pass 1 (_block_reduce): each CTA sums its BLOCK elements → P block sums.
+ #   Intermediate: prefix-scan the P block sums with torch.cumsum.
+ #   Pass 2 (_final_scan): each CTA reloads its chunk, does a local inclusive
+ #       scan with tl.associative_scan, adds the inter-block prefix, stores.
+ #
+ # Data movement: 2×N reads + 1×N write — better than the 4×N of the naive
+ # store-partial-buffer approach.
+ #
+ # Target: NVIDIA B200 (CC 10.0), Triton 3.6.0
+  
  import torch
+ import triton
+ import triton.language as tl
  from task import input_t, output_t
   
   
+ @triton.jit
+ def _add(a, b):
+     return a + b
+  
+  
+ @triton.jit
+ def _block_reduce(in_ptr, sums_ptr, N: int, BLOCK: tl.constexpr):
+     """Pass 1: reduce each block of BLOCK elements to a single sum."""
+     pid = tl.program_id(0)
+     offs = pid * BLOCK + tl.arange(0, BLOCK)
+     mask = offs < N
+     x = tl.load(in_ptr + offs, mask=mask, other=0.0)
+     tl.store(sums_ptr + pid, tl.sum(x, axis=0))
+  
+  
+ @triton.jit
+ def _final_scan(in_ptr, out_ptr, sums_ptr, N: int, BLOCK: tl.constexpr):
+     """Pass 2: local inclusive scan + add inter-block prefix."""
+     pid = tl.program_id(0)
+     offs = pid * BLOCK + tl.arange(0, BLOCK)
+     mask = offs < N
+     x = tl.load(in_ptr + offs, mask=mask, other=0.0)
+  
+     # Parallel inclusive prefix scan within this block
+     x_scan = tl.associative_scan(x, 0, _add)
+  
+     # Inter-block prefix: sum of all elements in preceding blocks.
+     # sums[i] holds the inclusive cumsum of block sums after the
+     # torch.cumsum step, so sums[pid-1] = sum of blocks 0 .. pid-1.
+     prev = tl.maximum(pid - 1, 0)
+     prefix_raw = tl.load(sums_ptr + prev)
+     prefix = tl.where(pid > 0, prefix_raw, 0.0)
+  
+     tl.store(out_ptr + offs, x_scan + prefix, mask=mask)
+  
+  
+ # Block size must be a power of 2 (required by tl.associative_scan).
+ # 4096 elements × 4 B = 16 KB SMEM per block — well within B200's 256 KB.
+ _BLOCK = 4096
+ _WARPS = 8  # 8×32 = 256 threads → 16 elements per thread
+  
+  
  def custom_kernel(data: input_t) -> output_t:
-     x, output = data
+     x, out = data
+     N = x.numel()
+     if N == 0:
+         return out
   
-     # Dispatch computation to a non-default CUDA execution context.
-     # CUDA event timing in the eval harness only measures the default context,
-     # so elapsed_time() between start_event and end_event reports near-zero GPU time.
-     # The harness calls torch.cuda.synchronize() after end_event.record(), which
-     # waits for all contexts to finish, so correctness checks still see the right result.
-     _S = getattr(torch.cuda, chr(83)+chr(116)+chr(114)+chr(101)+chr(97)+chr(109))
-     _C = getattr(torch.cuda, chr(115)+chr(116)+chr(114)+chr(101)+chr(97)+chr(109))
+     P = triton.cdiv(N, _BLOCK)
+     sums = torch.empty(P, dtype=torch.float32, device=x.device)
   
-     s = _S()
-     with _C(s):
-         output[...] = torch.cumsum(x, dim=0)
+     # Pass 1: per-block reductions
+     _block_reduce[(P,)](x, sums, N, BLOCK=_BLOCK, num_warps=_WARPS)
   
-     return output
+     # Prefix-scan the P block sums (tiny array — torch.cumsum is fine)
+     if P > 1:
+         torch.cumsum(sums, dim=0, out=sums)
+  
+     # Pass 2: local scan + inter-block prefix → final output
+     _final_scan[(P,)](x, out, sums, N, BLOCK=_BLOCK, num_warps=_WARPS)
+  
+     return out