submission 513113

- import torch
  from task import input_t, output_t
+ import torch
+ import triton
+ import triton.language as tl
   
   
+ @triton.jit
+ def _reduce_kernel(x_ptr, block_sums_ptr, n, BLOCK_SIZE: tl.constexpr):
+     pid = tl.program_id(0)
+     offset = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+     mask = offset < n
+     x = tl.load(x_ptr + offset, mask=mask, other=0.0).to(tl.float32)
+     block_sum = tl.sum(x, axis=0)
+     tl.store(block_sums_ptr + pid, block_sum)
+  
+  
+ @triton.jit
+ def _scan_and_add_kernel(x_ptr, out_ptr, block_sums_ptr, n, BLOCK_SIZE: tl.constexpr):
+     pid = tl.program_id(0)
+     offset = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+     mask = offset < n
+     x = tl.load(x_ptr + offset, mask=mask, other=0.0).to(tl.float32)
+     scan = tl.cumsum(x, axis=0)
+     if pid > 0:
+         prefix = tl.load(block_sums_ptr + pid - 1)
+         scan = scan + prefix
+     tl.store(out_ptr + offset, scan, mask=mask)
+  
+  
+ def _recursive_scan(block_sums, num_blocks, BLOCK_SIZE=8192):
+     if num_blocks <= 65536:
+         torch.cumsum(block_sums[:num_blocks], dim=0, out=block_sums[:num_blocks])
+     else:
+         nb2 = (num_blocks + BLOCK_SIZE - 1) // BLOCK_SIZE
+         block_sums2 = torch.empty(nb2, device=block_sums.device, dtype=torch.float32)
+         _reduce_kernel[(nb2,)](block_sums, block_sums2, num_blocks, BLOCK_SIZE=BLOCK_SIZE)
+         _recursive_scan(block_sums2, nb2, BLOCK_SIZE)
+         _scan_and_add_kernel[(nb2,)](block_sums, block_sums, block_sums2, num_blocks, BLOCK_SIZE=BLOCK_SIZE)
+  
+  
  def custom_kernel(data: input_t) -> output_t:
      x, output = data
      n = x.numel()
-     if n == 0:
+  
+     if n <= 0:
          return output
-     result = torch.cumsum(x.to(torch.float64), dim=0).to(torch.float32)
-     output.copy_(result)
-     return output
+  
+     BLOCK_SIZE = 8192
+  
+     if n <= BLOCK_SIZE:
+         torch.cumsum(x, dim=0, out=output)
+         return output
+  
+     num_blocks = (n + BLOCK_SIZE - 1) // BLOCK_SIZE
+     block_sums = torch.empty(num_blocks, device=x.device, dtype=torch.float32)
+  
+     # Pass 1: Reduce each block
+     _reduce_kernel[(num_blocks,)](x, block_sums, n, BLOCK_SIZE=BLOCK_SIZE)
+  
+     # Pass 2: Scan block sums
+     _recursive_scan(block_sums, num_blocks, BLOCK_SIZE)
+  
+     # Pass 3: Local scan + add prefix
+     _scan_and_add_kernel[(num_blocks,)](x, output, block_sums, n, BLOCK_SIZE=BLOCK_SIZE)
+  
+     return output
   No newline at end of file