submission 513378

  ⋯ 2 unchanged lines
  import triton.language as tl
  from task import input_t, output_t
   
- BLOCK_SIZE = 4096
- NUM_WARPS = 16
- BLOCK2 = 512  # enough for all benchmark sizes: max_blocks = ceil(52428800/4096) = 12800... need larger
   
- # Actually for BLOCK_SIZE=4096: ceil(52428800/4096) = 12800, need BLOCK2 >= 12800 -> 16384
- # For BLOCK_SIZE=16384: ceil(52428800/16384) = 3200, BLOCK2 = 4096
- BLOCK_SIZE = 16384
- NUM_WARPS = 32
- BLOCK2 = 4096  # 2^12 >= 3200
-  
  @triton.jit
- def reduction_pass1(input_ptr, partial_ptr, N, BLOCK_SIZE: tl.constexpr):
+ def _sum_kernel(
+     input_ptr,
+     output_ptr,
+     partial_ptr,
+     counter_ptr,
+     N,
+     NUM_PROGRAMS: tl.constexpr,
+     BLOCK_SIZE: tl.constexpr,
+     PARTIAL_BLOCK: tl.constexpr,
+ ):
      pid = tl.program_id(0)
-     offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
-     mask = offsets < N
-     x = tl.load(input_ptr + offsets, mask=mask, other=0.0).to(tl.float64)
-     tl.store(partial_ptr + pid, tl.sum(x, axis=0))
+     stride = NUM_PROGRAMS * BLOCK_SIZE
   
+     acc = tl.zeros([BLOCK_SIZE], dtype=tl.float64)
+     offset = pid * BLOCK_SIZE
+     while offset + BLOCK_SIZE <= N:
+         offsets = offset + tl.arange(0, BLOCK_SIZE)
+         vals = tl.load(input_ptr + offsets, eviction_policy='evict_first')
+         acc += vals.to(tl.float64)
+         offset += stride
+     if offset < N:
+         offsets = offset + tl.arange(0, BLOCK_SIZE)
+         mask = offsets < N
+         vals = tl.load(input_ptr + offsets, mask=mask, other=0.0, eviction_policy='evict_first')
+         acc += vals.to(tl.float64)
   
- @triton.jit
- def reduction_pass2(partial_ptr, output_ptr, M, BLOCK2: tl.constexpr):
-     offsets = tl.arange(0, BLOCK2)
-     mask = offsets < M
-     x = tl.load(partial_ptr + offsets, mask=mask, other=0.0)
-     tl.store(output_ptr, tl.sum(x, axis=0).to(tl.float32))
+     partial_sum = tl.sum(acc, axis=0)
+     tl.store(partial_ptr + pid, partial_sum)
   
+     count = tl.atomic_add(counter_ptr, 1, sem="release")
   
- _partial_buf = None
- _warmed_up = False
+     if count == NUM_PROGRAMS - 1:
+         tl.atomic_add(counter_ptr, 0, sem="acquire")
+         p_offsets = tl.arange(0, PARTIAL_BLOCK)
+         p_mask = p_offsets < NUM_PROGRAMS
+         partials = tl.load(partial_ptr + p_offsets, mask=p_mask, other=0.0)
+         total = tl.sum(partials, axis=0)
+         tl.store(output_ptr, total.to(tl.float32))
+         tl.store(counter_ptr, 0)
   
- def _ensure_buf(n_blocks):
-     global _partial_buf
-     max_needed = 4096  # enough for BLOCK_SIZE=16384 up to 52M elements
-     if _partial_buf is None:
-         _partial_buf = torch.empty(max_needed, device='cuda', dtype=torch.float64)
   
+ _buffers = {}
   
- def _warmup():
-     global _warmed_up
-     if _warmed_up:
-         return
-     _ensure_buf(4096)
-     # Warmup pass1
-     dummy_in = torch.zeros(BLOCK_SIZE, device='cuda', dtype=torch.float32)
-     dummy_out = torch.zeros(1, device='cuda', dtype=torch.float32)
-     reduction_pass1[(1,)](dummy_in, _partial_buf, BLOCK_SIZE, BLOCK_SIZE=BLOCK_SIZE, num_warps=NUM_WARPS)
-     reduction_pass2[(1,)](_partial_buf, dummy_out, 1, BLOCK2=BLOCK2, num_warps=16)
-     torch.cuda.synchronize()
-     _warmed_up = True
+ def _get_buffers(device, num_programs):
+     key = (device, num_programs)
+     if key not in _buffers:
+         _buffers[key] = (
+             torch.empty(num_programs, device=device, dtype=torch.float64),
+             torch.zeros(1, device=device, dtype=torch.int32),
+         )
+     return _buffers[key]
   
+ def _next_pow2(x):
+     p = 1
+     while p < x:
+         p *= 2
+     return p
   
- _warmup()
-  
-  
  def custom_kernel(data: input_t) -> output_t:
      input_tensor, output_tensor = data
      N = input_tensor.numel()
-     num_blocks = triton.cdiv(N, BLOCK_SIZE)
-     
-     _ensure_buf(num_blocks)
-     partial = _partial_buf[:num_blocks]
-     
-     reduction_pass1[(num_blocks,)](
-         input_tensor, partial, N,
+     device = input_tensor.device
+  
+     if N <= 4096:
+         num_programs = min(16, max(1, (N + 255) // 256))
+         BLOCK_SIZE = 1024
+         nw = 4
+         ns = 2
+     elif N <= 65536:
+         num_programs = min(64, (N + 1023) // 1024)
+         BLOCK_SIZE = 2048
+         nw = 8
+         ns = 2
+     elif N <= 1048576:
+         num_programs = 128
+         BLOCK_SIZE = 4096
+         nw = 8
+         ns = 4
+     elif N <= 6553600:
+         num_programs = 256
+         BLOCK_SIZE = 4096
+         nw = 8
+         ns = 4
+     else:
+         num_programs = 512
+         BLOCK_SIZE = 4096
+         nw = 8
+         ns = 4
+  
+     PARTIAL_BLOCK = _next_pow2(num_programs)
+     partial_sums, counter = _get_buffers(device, num_programs)
+  
+     _sum_kernel[(num_programs,)](
+         input_tensor,
+         output_ptr=output_tensor,
+         partial_ptr=partial_sums,
+         counter_ptr=counter,
+         N=N,
+         NUM_PROGRAMS=num_programs,
          BLOCK_SIZE=BLOCK_SIZE,
-         num_warps=NUM_WARPS,
+         PARTIAL_BLOCK=PARTIAL_BLOCK,
+         num_warps=nw,
+         num_stages=ns,
      )
-     
-     reduction_pass2[(1,)](
-         partial, output_tensor, num_blocks,
-         BLOCK2=BLOCK2,
-         num_warps=16,
-     )
-     
-     return output_tensor[0]
   No newline at end of file
+  
+     return output_tensor.squeeze()
   No newline at end of file