submission 774931

  import torch
+ import triton
+ import triton.language as tl
  from task import input_t, output_t
   
+ # A100 optimized kernel - do not remove
+ # Note: Running on A100 (108 SMs, SM80)
+ #
+ # Why Triton beats current pure-PyTorch kernel:
+ #   Current:  sum(dtype=float64)  → fp64 scalar (alloc #1)
+ #             .to(float32)        → fp32 scalar (alloc #2, separate GPU kernel)
+ #   This:     pass1 → 512 fp64 partials into pre-allocated _scratch (no new alloc)
+ #             pass2 → fp32 result fused directly into pre-allocated output_tensor
+ #   Savings:  2 fewer temp tensor allocations, 1 fewer GPU kernel launch,
+ #             fp64→fp32 cast fused into the store in pass2.
+ #
+ # Why previous Triton submissions failed (exit 112 = check_implementation mismatch):
+ #   They returned output_tensor with shape (1,) — ref_kernel returns shape ().
+ #   Fix: return output_tensor.view([]) — a 0-dim scalar view of the same storage.
+ #
+ # JIT warmup: the benchmark harness calls run_single_benchmark on tests[0] before
+ # timing begins. Triton compiles once there. All timed calls hit the kernel cache.
+ # The test phase (5 small cases, same spawn worker) also compiles once on first call.
+  
+ GRID = 512    # 2^9 — power of 2 required by tl.arange in pass2
+ BLOCK = 4096  # 2^12 — tile size; 16KB/tile saturates A100 HBM at 512 blocks
+  
+  
+ @triton.jit
+ def _pass1(data_ptr, partial_ptr, N, BLOCK: tl.constexpr, GRID: tl.constexpr):
+     """
+     512 programs, each grid-strides through its slice of fp32 input.
+     Accumulates in fp64 registers — no intermediate fp64 tensor written to HBM.
+     Writes one fp64 partial sum per program to _scratch.
+     """
+     pid = tl.program_id(0)
+     acc = tl.zeros((BLOCK,), dtype=tl.float64)
+     start = pid * BLOCK
+     stride = GRID * BLOCK
+     while start < N:
+         offs = start + tl.arange(0, BLOCK)
+         mask = offs < N
+         acc += tl.load(data_ptr + offs, mask=mask, other=0.0).to(tl.float64)
+         start += stride
+     tl.store(partial_ptr + pid, tl.sum(acc, axis=0))
+  
+  
+ @triton.jit
+ def _pass2(partial_ptr, out_ptr, GRID: tl.constexpr):
+     """
+     Single program sums 512 fp64 partials and stores fp32 result directly
+     into the pre-allocated output_tensor — no separate cast kernel.
+     """
+     offs = tl.arange(0, GRID)            # GRID must be power of 2
+     total = tl.sum(tl.load(partial_ptr + offs), axis=0)
+     tl.store(out_ptr, total.to(tl.float32))  # fp64→fp32 fused here
+  
+  
+ # Module-level scratch: allocated lazily on first call so module import
+ # doesn't require CUDA to be ready (safe for the spawn pool worker).
+ _scratch: torch.Tensor | None = None
+  
+  
  def custom_kernel(data: input_t) -> output_t:
+     global _scratch
      input_tensor, output_tensor = data
-     # .to(float64).sum() allocates a full fp64 tensor in HBM then sums it:
-     #   read fp32 -> write fp64 (new alloc) -> read fp64 -> scalar
-     # .sum(dtype=float64) fuses the cast into the reduction kernel:
-     #   read fp32 -> scalar  (no intermediate tensor)
-     # Returns a scalar (shape []) matching the reference's return type.
-     return input_tensor.sum(dtype=torch.float64).to(torch.float32)
   No newline at end of file
+     N = input_tensor.numel()
+  
+     if _scratch is None:
+         _scratch = torch.empty(GRID, device="cuda", dtype=torch.float64)
+  
+     _pass1[(GRID,)](
+         input_tensor, _scratch, N,
+         BLOCK=BLOCK, GRID=GRID,
+         num_warps=8,
+     )
+     _pass2[(1,)](
+         _scratch, output_tensor,
+         GRID=GRID,
+         num_warps=4,
+     )
+  
+     # output_tensor is shape (1,); ref_kernel returns shape ().
+     # view([]) gives a 0-dim scalar backed by the same storage — no copy.
+     return output_tensor.view([])
   No newline at end of file