submission 68268

import torch
import triton
import triton.language as tl

@triton.jit
def sum_kernel_stage1(
    input_ptr,
    partial_sums_ptr,
    n_elements,
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    
    data = tl.load(input_ptr + offsets, mask=mask, other=0.0)
    block_sum = tl.sum(data)
    tl.store(partial_sums_ptr + pid, block_sum)

@triton.jit
def sum_kernel_stage2(
    partial_sums_ptr,
    partial_sums2_ptr,
    n_partials,
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_partials
    
    data = tl.load(partial_sums_ptr + offsets, mask=mask, other=0.0)
    block_sum = tl.sum(data)
    tl.store(partial_sums2_ptr + pid, block_sum)

def custom_kernel(data):
    input_tensor, output_tensor = data
    n_elements = input_tensor.numel()
    
    # 安定性とパフォーマンスのバランス
    BLOCK_SIZE_1 = 4096
    n_blocks_1 = triton.cdiv(n_elements, BLOCK_SIZE_1)
    partial_sums_1 = torch.empty(n_blocks_1, device='cuda', dtype=torch.float32)
    
    # num_warps=8を明示的に指定（BLOCK_SIZE=4096に最適）
    sum_kernel_stage1[(n_blocks_1,)](
        input_tensor,
        partial_sums_1,
        n_elements,
        BLOCK_SIZE=BLOCK_SIZE_1,
        num_warps=16,
    )
    
    # Stage 2
    BLOCK_SIZE_2 = 1024
    n_blocks_2 = triton.cdiv(n_blocks_1, BLOCK_SIZE_2)
    
    if n_blocks_2 > 1:
        partial_sums_2 = torch.empty(n_blocks_2, device='cuda', dtype=torch.float32)
        sum_kernel_stage2[(n_blocks_2,)](
            partial_sums_1,
            partial_sums_2,
            n_blocks_1,
            BLOCK_SIZE=BLOCK_SIZE_2,
            num_warps=16,
        )
        # result = partial_sums_2.to(torch.float64).sum().to(torch.float32)
        result = partial_sums_2.sum()
    else:
        # result = partial_sums_1.to(torch.float64).sum().to(torch.float32)
        result = partial_sums_1.sum()
    
    return result