submission 131823

  ⋯ 2 unchanged lines
  import triton.language as tl
  from task import input_t, output_t
  
- # Autotuneは廃止（オーバーヘッド完全除去 & Atomicの安全性確保のため）
- # 以前の計測で、巨大ブロックには warps=32, stages=4 が効くことが分かっています
+ # BLOCK_SIZEは固定しますが、内部の実行パラメータを徹底的にチューニングします
+ # pre_hook不要（Storeは上書きなので何度実行しても安全）
+ @triton.autotune(
+     configs=[
+         # BLOCK_SIZE=32768 に対する最適設定を探る
+         # Warpsを増やすことで、大量の要素処理時のレイテンシを隠蔽
+         triton.Config({}, num_warps=16, num_stages=2),
+         triton.Config({}, num_warps=32, num_stages=2),
+         triton.Config({}, num_warps=16, num_stages=4),
+         triton.Config({}, num_warps=32, num_stages=4),
+         # 環境によってはWarp数が多すぎるとレジスタ溢れするため、少なめの設定も保険に入れる
+         triton.Config({}, num_warps=8, num_stages=2),
+     ],
+     key=['n_elements'],
+ )
  @triton.jit
- def _sum_kernel_atomic_fixed(
+ def _sum_kernel_map_fixed_block(
      x_ptr,
-     y_ptr,
+     temp_ptr,
      n_elements,
+     # BLOCK_SIZEはコンパイル時定数として渡すが、値は固定
      BLOCK_SIZE: tl.constexpr,
  ):
      pid = tl.program_id(axis=0)
  ⋯ 4 unchanged lines
      mask = offsets < n_elements
  
      # Load & Cast
-     # 32768要素を一括ロード
+     # 32768要素を一気にロード。Tritonが自動でベクトル化・分割して最適化します
      x = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
  
      # ブロック内リダクション
      block_sum = tl.sum(x, axis=0)
  
-     # Atomic Add
-     # ブロックサイズが32768と巨大なため、競合回数は極小（約1500回）
-     # Autotuneを使わないため、カーネルは「必ず1回」しか走らない。
-     # したがって、値が累積するバグやnull落ちのリスクがない。
-     tl.atomic_add(y_ptr, block_sum)
+     # Store (Atomicを使わず上書き)
+     # これにより "null" 落ちやテスト失敗（非決定性）を回避
+     tl.store(temp_ptr + pid, block_sum)
  
  def custom_kernel(data: input_t) -> output_t:
      """
-     Triton Single-Kernel Atomic (Fixed Config).
-     - Removes the overhead of the second kernel launch (torch.sum).
-     - Hardcoded parameters ensure stability and zero python overhead.
+     Triton Optimized Map-Reduce (Fixed Huge Block).
+     - Fixes BLOCK_SIZE to Maximize Bandwidth & Determine Grid Size.
+     - Autotunes Warps/Stages for Hardware Optimization.
+     - Uses torch.empty for zero-overhead allocation.
      """
      input_tensor, _ = data
-     
-     # 【極限チューニング】
-     # ベンチマークの入力が常にContiguousだと分かっているなら、
-     # このチェックと変換を外すことで さらに 1-2µs 削れます。
-     # 安全性のため残していますが、コメントアウトすれば最速です。
      x = input_tensor.view(-1)
-     # if not x.is_contiguous():
-     #     x = x.contiguous()
+     
+     if not x.is_contiguous():
+         x = x.contiguous()
          
      n_elements = x.numel()
      
-     # 固定パラメータ (Autotuneの時間を待たず、即座に最適設定で走らせる)
+     # 【高速化の鍵】BLOCK_SIZEを32768に固定
+     # 5000万要素 ÷ 32768 ≒ 1526 ブロック
+     # これによりAtomic競合を無くしつつ、ループオーバーヘッドも最小化
      BLOCK_SIZE = 32768
      
-     # 出力バッファ
-     # atomic_addを使うため、zerosで初期化必須
-     # 1要素のzerosは一瞬で終わるため、2回目のカーネル起動(torch.sum)より圧倒的に速い
-     output = torch.zeros(1, device=input_tensor.device, dtype=torch.float32)
+     # グリッドサイズを確定
+     grid_size = triton.cdiv(n_elements, BLOCK_SIZE)
      
-     # Grid計算
-     grid = (triton.cdiv(n_elements, BLOCK_SIZE), )
+     # 【高速化の鍵】torch.emptyを使用
+     # グリッドサイズ分だけ確保。初期化しない（カーネルが全要素を上書きするため安全）
+     # zerosの初期化コスト（数µs）をカット
+     temp_buffer = torch.empty(grid_size, device=input_tensor.device, dtype=torch.float32)
      
-     # カーネル実行
-     # num_warps=32: ハイエンドGPU向けの高並列設定
-     # num_stages=4: メモリレイテンシ隠蔽設定
-     _sum_kernel_atomic_fixed[grid](
+     # Grid定義
+     grid = (grid_size, )
+     
+     # カーネル実行（内部でAutotuneが走り、最適なnum_warpsが選ばれる）
+     # kwargsでBLOCK_SIZEを渡す
+     _sum_kernel_map_fixed_block[grid](
          x, 
-         output, 
+         temp_buffer, 
          n_elements, 
-         BLOCK_SIZE=BLOCK_SIZE,
-         num_warps=16,#16
-         num_stages=4,#4
+         BLOCK_SIZE=BLOCK_SIZE
      )
      
-     return output[0]
   No newline at end of file
+     # Phase 2: わずか~1500要素の足し算
+     # ここはPyTorchのC++実装が一瞬で処理する
+     return torch.sum(temp_buffer)
   No newline at end of file