submission 131828

import torch
import triton
import triton.language as tl
from task import input_t, output_t

# 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_map_fixed_block(
    x_ptr,
    temp_ptr,
    n_elements,
    # BLOCK_SIZEはコンパイル時定数として渡すが、値は固定
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(axis=0)
    
    # 担当領域計算
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements

    # Load & Cast
    # 32768要素を一気にロード。Tritonが自動でベクトル化・分割して最適化します
    x = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float32)

    # ブロック内リダクション
    block_sum = tl.sum(x, axis=0)

    # Store (Atomicを使わず上書き)
    # これにより "null" 落ちやテスト失敗（非決定性）を回避
    tl.store(temp_ptr + pid, block_sum)

def custom_kernel(data: input_t) -> output_t:
    """
    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
    x = input_tensor.view(-1)
    
    if not x.is_contiguous():
        x = x.contiguous()
        
    n_elements = x.numel()
    
    # 【高速化の鍵】BLOCK_SIZEを32768に固定
    # 5000万要素 ÷ 32768 ≒ 1526 ブロック
    # これによりAtomic競合を無くしつつ、ループオーバーヘッドも最小化
    BLOCK_SIZE = 32768
    
    # グリッドサイズを確定
    grid_size = triton.cdiv(n_elements, BLOCK_SIZE)
    
    # 【高速化の鍵】torch.emptyを使用
    # グリッドサイズ分だけ確保。初期化しない（カーネルが全要素を上書きするため安全）
    # zerosの初期化コスト（数µs）をカット
    temp_buffer = torch.empty(grid_size, device=input_tensor.device, dtype=torch.float32)
    
    # Grid定義
    grid = (grid_size, )
    
    # カーネル実行（内部でAutotuneが走り、最適なnum_warpsが選ばれる）
    # kwargsでBLOCK_SIZEを渡す
    _sum_kernel_map_fixed_block[grid](
        x, 
        temp_buffer, 
        n_elements, 
        BLOCK_SIZE=BLOCK_SIZE
    )
    
    # Phase 2: わずか~1500要素の足し算
    # ここはPyTorchのC++実装が一瞬で処理する
    return torch.sum(temp_buffer)