submission 545121

  ⋯ 1 unchanged lines
  os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
   
  import torch
+ import triton
+ import triton.language as tl
  from task import input_t, output_t
   
+ # Two-level reduction: GPU blocks reduce locally, then CPU sums partials
+ # Use float64 accumulation for accuracy matching reference
+  
+ @triton.jit
+ def sum_reduce_kernel(
+     x_ptr, partial_ptr, n_elements,
+     BLOCK_SIZE: tl.constexpr,
+ ):
+     pid = tl.program_id(0)
+     offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+     mask = offsets < n_elements
+     # Load as float32, accumulate in float64
+     x = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float64)
+     block_sum = tl.sum(x, axis=0)
+     tl.store(partial_ptr + pid, block_sum)
+  
+ @triton.jit
+ def sum_final_kernel(
+     partial_ptr, out_ptr, n_partials,
+     BLOCK_SIZE: tl.constexpr,
+ ):
+     offsets = tl.arange(0, BLOCK_SIZE)
+     mask = offsets < n_partials
+     x = tl.load(partial_ptr + offsets, mask=mask, other=0.0)
+     total = tl.sum(x, axis=0)
+     tl.store(out_ptr, total.to(tl.float32))
+  
  def custom_kernel(data: input_t) -> output_t:
      data, output = data
-     # Reference does: data.to(torch.float64).sum().to(torch.float32)
-     # and returns a scalar. We need to match that exactly.
-     result = data.to(torch.float64).sum().to(torch.float32)
-     return result
+     n = data.numel()
+     BLOCK_SIZE = 4096
+     n_blocks = (n + BLOCK_SIZE - 1) // BLOCK_SIZE
+     partial = torch.empty(n_blocks, device=data.device, dtype=torch.float64)
+     sum_reduce_kernel[(n_blocks,)](data, partial, n, BLOCK_SIZE=BLOCK_SIZE)
+     
+     # Second level reduction on GPU if too many partials
+     if n_blocks <= 65536:
+         # Use power-of-2 BLOCK_SIZE for final reduction
+         FINAL_BS = 1
+         while FINAL_BS < n_blocks:
+             FINAL_BS *= 2
+         if FINAL_BS > 65536:
+             FINAL_BS = 65536
+         out_scalar = torch.empty(1, device=data.device, dtype=torch.float32)
+         sum_final_kernel[(1,)](partial, out_scalar, n_blocks, BLOCK_SIZE=FINAL_BS)
+         return out_scalar[0]
+     else:
+         return partial.sum().to(torch.float32)