submission 774965

  ⋯ 1 unchanged lines
  import triton
  import triton.language as tl
  from task import input_t, output_t
+ GRID = 512
+ BLOCK = 4096
   
- # 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
  ⋯ 8 unchanged lines
   
  @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
+     offs = tl.arange(0, GRID)
      total = tl.sum(tl.load(partial_ptr + offs), axis=0)
-     tl.store(out_ptr, total.to(tl.float32))  # fp64→fp32 fused here
+     tl.store(out_ptr, total.to(tl.float32))
   
   
- # 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
   
   
  ⋯ 6 unchanged lines
          _scratch = torch.empty(GRID, device="cuda", dtype=torch.float64)
   
      _pass1[(GRID,)](
-         input_tensor, _scratch, N,
-         BLOCK=BLOCK, GRID=GRID,
+         input_tensor,
+         _scratch,
+         N,
+         BLOCK=BLOCK,
+         GRID=GRID,
          num_warps=8,
      )
      _pass2[(1,)](
-         _scratch, output_tensor,
+         _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