submission 512006

  ⋯ 12 unchanged lines
  from task import input_t, output_t
   
   
- BLOCK_SIZE = 1024
- DEFAULT_NUM_WARPS = 8
+ FIRST_BLOCK_SIZE = 2048
+ REDUCE_BLOCK_SIZE = 1024
+ FIRST_PASS_WARPS = 8
+ LATE_PASS_WARPS = 4
   
- _PARTIAL_CACHE_A = {}
- _PARTIAL_CACHE_B = {}
+ _BUF_A = {}
+ _BUF_B = {}
   
   
  if _TRITON_AVAILABLE:
   
      @triton.jit
-     def _partial_sum_kernel(
+     def _reduce_block_kernel(
          x_ptr,
          partial_ptr,
          n_elements,
  ⋯ 2 unchanged lines
          pid = tl.program_id(axis=0)
          offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
          mask = offsets < n_elements
+         vals = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float64)
+         acc = tl.sum(vals, axis=0)
+         tl.store(partial_ptr + pid, acc)
   
-         values = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float64)
-         partial = tl.sum(values, axis=0)
   
-         tl.store(partial_ptr + pid, partial)
+ def _get_buffer(cache: dict, device: torch.device, needed: int) -> torch.Tensor:
+     key = (device.type, device.index)
+     buf = cache.get(key)
+     if buf is None or buf.numel() < needed:
+         alloc = 1 << (max(1, needed) - 1).bit_length()
+         buf = torch.empty((alloc,), device=device, dtype=torch.float64)
+         cache[key] = buf
+     return buf[:needed]
   
   
- def _get_partial_buffer(cache: dict, device: torch.device, num_partials: int) -> torch.Tensor:
-     key = (device.type, device.index)
-     cached = cache.get(key)
-     if cached is None or cached.numel() < num_partials:
-         alloc_size = 1 << (max(1, num_partials) - 1).bit_length()
-         cached = torch.empty((alloc_size,), device=device, dtype=torch.float64)
-         cache[key] = cached
-     return cached[:num_partials]
+ def _launch_reduce_block(
+     inp: torch.Tensor,
+     out: torch.Tensor,
+     n_elements: int,
+     block_size: int,
+     num_warps: int,
+ ) -> None:
+     grid = (out.numel(),)
+     _reduce_block_kernel[grid](
+         inp,
+         out,
+         n_elements,
+         BLOCK_SIZE=block_size,
+         num_warps=num_warps,
+         num_stages=2,
+     )
   
   
- def _triton_reduce_sum(x: torch.Tensor) -> torch.Tensor:
-     current = x
-     n_elements = current.numel()
-     use_a = True
+ def _triton_sum_fp64(x: torch.Tensor) -> torch.Tensor:
+     n0 = x.numel()
+     first_block = FIRST_BLOCK_SIZE if n0 >= (1 << 24) else REDUCE_BLOCK_SIZE
+     b1 = triton.cdiv(n0, first_block)
+     p1 = _get_buffer(_BUF_A, x.device, b1)
+     _launch_reduce_block(x, p1, n0, first_block, FIRST_PASS_WARPS)
+     if b1 == 1:
+         return p1[0]
   
-     num_warps = DEFAULT_NUM_WARPS
-     if n_elements < (1 << 18):
-         num_warps = 4
+     b2 = triton.cdiv(b1, REDUCE_BLOCK_SIZE)
+     p2 = _get_buffer(_BUF_B, x.device, b2)
+     _launch_reduce_block(p1, p2, b1, REDUCE_BLOCK_SIZE, LATE_PASS_WARPS)
+     if b2 == 1:
+         return p2[0]
   
-     while True:
-         num_blocks = triton.cdiv(n_elements, BLOCK_SIZE)
+     b3 = triton.cdiv(b2, REDUCE_BLOCK_SIZE)
+     p3 = _get_buffer(_BUF_A, x.device, b3)
+     _launch_reduce_block(p2, p3, b2, REDUCE_BLOCK_SIZE, LATE_PASS_WARPS)
+     if b3 == 1:
+         return p3[0]
   
+     current = p3
+     n_current = b3
+     use_a = False
+     while n_current > 1:
+         b = triton.cdiv(n_current, REDUCE_BLOCK_SIZE)
          if use_a:
-             partials = _get_partial_buffer(_PARTIAL_CACHE_A, x.device, num_blocks)
+             nxt = _get_buffer(_BUF_A, x.device, b)
          else:
-             partials = _get_partial_buffer(_PARTIAL_CACHE_B, x.device, num_blocks)
-  
-         _partial_sum_kernel[(num_blocks,)](
-             current,
-             partials,
-             n_elements,
-             BLOCK_SIZE=BLOCK_SIZE,
-             num_warps=num_warps,
-             num_stages=2,
-         )
-  
-         if num_blocks == 1:
-             return partials[0]
-  
-         current = partials
-         n_elements = num_blocks
+             nxt = _get_buffer(_BUF_B, x.device, b)
+         _launch_reduce_block(current, nxt, n_current, REDUCE_BLOCK_SIZE, LATE_PASS_WARPS)
+         current = nxt
+         n_current = b
          use_a = not use_a
-         num_warps = 4
+     return current[0]
   
   
  def custom_kernel(data: input_t) -> output_t:
-     x, _output = data
+     x, output = data
      if not x.is_contiguous():
          x = x.contiguous()
   
      if x.numel() == 0:
-         total_fp64 = torch.zeros((), device=x.device, dtype=torch.float64)
+         total = torch.zeros((), device=x.device, dtype=torch.float64)
      elif x.is_cuda and _TRITON_AVAILABLE:
          try:
-             total_fp64 = _triton_reduce_sum(x)
+             total = _triton_sum_fp64(x)
          except Exception:
-             total_fp64 = x.to(torch.float64).sum()
+             total = x.to(torch.float64).sum()
      else:
-         total_fp64 = x.to(torch.float64).sum()
+         total = x.to(torch.float64).sum()
   
-     return total_fp64.to(torch.float32)
+     out_scalar = output.view(())
+     out_scalar.copy_(total)
+     return out_scalar