submission 233810

  import torch
- import triton
- import triton.language as tl
+ from torch.utils.cpp_extension import load_inline
  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
+ # =============================================================================
+ # 実験22: 2アキュムレータ + __threadfence()最小化
+ # 目標: 平均≤46.651μs
+ # 戦略:
+ #   - 実験21の2アキュムレータILP（効果確認済み）
+ #   - __threadfence()を最後のブロックのみに移動（実験18方式）
+ #   - 実験18では2要素並列ロードが悪影響だったが、
+ #     2アキュムレータなら安定性を維持できる可能性
+ #
+ # 期待効果:
+ #   - __threadfence() 1600回 → 1回で1-2μs削減
+ #   - 2アキュムレータによる安定性維持
+ # =============================================================================
  
-     # Load & Cast
-     # 32768要素を一気にロード。Tritonが自動でベクトル化・分割して最適化します
-     x = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+ cuda_src = """
+ #include <cuda_runtime.h>
  
-     # ブロック内リダクション
-     block_sum = tl.sum(x, axis=0)
+ // =============================================================================
+ // Warp-level reduction using shuffle (決定論的順序)
+ // =============================================================================
+ __device__ __forceinline__ float warpReduceSum(float val) {
+     #pragma unroll
+     for (int offset = 16; offset > 0; offset /= 2) {
+         val += __shfl_down_sync(0xffffffff, val, offset);
+     }
+     return val;
+ }
  
-     # Store (Atomicを使わず上書き)
-     # これにより "null" 落ちやテスト失敗（非決定性）を回避
-     tl.store(temp_ptr + pid, block_sum)
+ // =============================================================================
+ // Warp-level reduction for double precision
+ // =============================================================================
+ __device__ __forceinline__ double warpReduceSumDouble(double val) {
+     #pragma unroll
+     for (int offset = 16; offset > 0; offset /= 2) {
+         val += __shfl_down_sync(0xffffffff, val, offset);
+     }
+     return val;
+ }
  
+ // =============================================================================
+ // Block-level reduction using shared memory (決定論的順序)
+ // =============================================================================
+ __device__ __forceinline__ float blockReduceSum(float val, int numWarps) {
+     __shared__ float shared[32];  // 1 value per warp (max 32 warps)
+     int lane = threadIdx.x & 31;  // threadIdx.x % 32
+     int wid = threadIdx.x >> 5;   // threadIdx.x / 32
+ 
+     // Warp内リダクション
+     val = warpReduceSum(val);
+ 
+     // 各Warpの結果をshared memoryに書き込み
+     if (lane == 0) shared[wid] = val;
+     __syncthreads();
+ 
+     // 最初のWarpが全Warpの結果をリダクション
+     val = (threadIdx.x < numWarps) ? shared[lane] : 0.0f;
+     if (wid == 0) val = warpReduceSum(val);
+ 
+     return val;
+ }
+ 
+ // =============================================================================
+ // 静的バッファ（Phase1→Phase2の中間結果）
+ // =============================================================================
+ __device__ float g_partial_sums[2048];        // 部分和バッファ
+ __device__ unsigned int g_completed_count;    // 完了ブロックカウンタ
+ 
+ // =============================================================================
+ // Single-Kernel Reduction with 2-Accumulator + Minimal Fence
+ // =============================================================================
+ __global__ void reduce_single_kernel(
+     const float4* __restrict__ input4,
+     const float* __restrict__ input,
+     float* __restrict__ output,
+     const int n_elements,
+     const int n_vec4
+ ) {
+     // =========================================================================
+     // Phase 1: Main reduction with 2 accumulators
+     // =========================================================================
+     float sum0 = 0.0f;
+     float sum1 = 0.0f;
+ 
+     int tid = blockIdx.x * blockDim.x + threadIdx.x;
+     int stride = blockDim.x * gridDim.x;
+ 
+     // Grid-stride loop with 2 accumulators for ILP
+     for (int i = tid; i < n_vec4; i += stride) {
+         float4 v = input4[i];
+         sum0 += v.x + v.z;
+         sum1 += v.y + v.w;
+     }
+ 
+     // 残りの要素（n_elements % 4 != 0 の場合）
+     int tail_start = n_vec4 * 4;
+     for (int j = tail_start + tid; j < n_elements; j += stride) {
+         sum0 += input[j];
+     }
+ 
+     // 2つのアキュムレータを統合
+     float sum = sum0 + sum1;
+ 
+     // Block内リダクション (8 warps = 256 threads)
+     sum = blockReduceSum(sum, 8);
+ 
+     // =========================================================================
+     // Phase 1 → Phase 2 遷移（__threadfence()なし）
+     // =========================================================================
+     __shared__ bool is_last_block;
+ 
+     if (threadIdx.x == 0) {
+         // 部分和を書き込み（fenceなし）
+         g_partial_sums[blockIdx.x] = sum;
+ 
+         // 完了カウンタをインクリメント
+         unsigned int completed = atomicAdd(&g_completed_count, 1);
+ 
+         // 最後のブロックかどうかを判定
+         is_last_block = (completed == gridDim.x - 1);
+     }
+ 
+     __syncthreads();
+ 
+     // =========================================================================
+     // Phase 2: Final aggregation (最後のブロックのみ実行)
+     // =========================================================================
+     if (is_last_block) {
+         // 最後のブロックのみ__threadfence()を呼び出し
+         __threadfence();
+ 
+         // 1600要素をdouble精度で集約
+         double final_sum = 0.0;
+ 
+         for (int k = threadIdx.x; k < gridDim.x; k += blockDim.x) {
+             final_sum += (double)g_partial_sums[k];
+         }
+ 
+         // Warp内リダクション (double精度)
+         final_sum = warpReduceSumDouble(final_sum);
+ 
+         // Warp間リダクション (8 warps)
+         __shared__ double shared_double[8];
+         int lane = threadIdx.x & 31;
+         int wid = threadIdx.x >> 5;
+ 
+         if (lane == 0) shared_double[wid] = final_sum;
+         __syncthreads();
+ 
+         // 最初のWarpが最終集約
+         if (wid == 0) {
+             final_sum = (threadIdx.x < 8) ? shared_double[lane] : 0.0;
+             final_sum = warpReduceSumDouble(final_sum);
+ 
+             if (threadIdx.x == 0) {
+                 *output = (float)final_sum;
+ 
+                 // 次回実行のためにカウンタをリセット
+                 g_completed_count = 0;
+             }
+         }
+     }
+ }
+ 
+ // =============================================================================
+ // Host API
+ // =============================================================================
+ torch::Tensor vector_sum(torch::Tensor input, torch::Tensor output) {
+     const int n_elements = input.numel();
+     const int n_vec4 = n_elements / 4;
+ 
+     const int threads = 256;
+     const int blocks = 1600;
+ 
+     reduce_single_kernel<<<blocks, threads>>>(
+         reinterpret_cast<const float4*>(input.data_ptr<float>()),
+         input.data_ptr<float>(),
+         output.data_ptr<float>(),
+         n_elements,
+         n_vec4
+     );
+ 
+     return output;
+ }
+ """
+ 
+ cpp_src = """
+ #include <torch/extension.h>
+ 
+ torch::Tensor vector_sum(torch::Tensor input, torch::Tensor output);
+ """
+ 
+ _module = None
+ _output_buffer = None
+ 
+ def _get_module():
+     global _module
+     if _module is None:
+         _module = load_inline(
+             name='vectorsum_cuda_v22',
+             cuda_sources=[cuda_src],
+             cpp_sources=[cpp_src],
+             functions=['vector_sum'],
+             extra_cuda_cflags=['-O3', '--use_fast_math', '-lineinfo'],
+             verbose=False
+         )
+     return _module
+ 
+ 
  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.
+     実験22: 2アキュムレータ + __threadfence()最小化.
+ 
+     設計方針:
+     - 実験21の2アキュムレータを維持
+     - __threadfence()を最後のブロックのみに移動
+     - 1600回→1回でオーバーヘッド削減
+ 
+     期待性能: 48-49μs（目標46.651μs）
      """
+     global _output_buffer
      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)
   No newline at end of file
+ 
+     if _output_buffer is None or _output_buffer.device != input_tensor.device:
+         _output_buffer = torch.empty(1, device=input_tensor.device, dtype=torch.float32)
+ 
+     module = _get_module()
+     module.vector_sum(input_tensor.view(-1), _output_buffer)
+ 
+     return _output_buffer[0]