submission 68221 · KernelIndex

Use it

Vendorable · source mirrored · license unknownView source →

No package. Vendor the mirrored source: 72 lines, June 9 Researcher Reciprocity License v1.0.

all_triton.py

curl "https://kernelindex.com/api/v1/implementations/kernelbot-vectorsum-v2-68221?include=source"

interfacepython

Compatibility

measured onNVIDIA A100

declared hardwareNVIDIA A100

architecturessm_80

dtypesfp32

Benchmark evidence

1 measurement across 1 GPU, fastest first.

Operation / workload

Hardware

Latency

Rank

Observed

Vector sum reductionsuite of 6 cases

NVIDIA A100

152.7µs

#60 of 96

2025-11-08

Run →

Reported · How evidence levels are derived →

Source and license

sourceavailable

revision digestsha256:e973b130825bd91c13554804822a2c9ab18789edfeb5d3de0a36ccb69e36cc4d

license declaredunknown

license concludedunknown

authorskoshibat

imported2026-08-15

Kernel source

all_triton.py72 lines

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,
):
    """Stage 2: 部分和をさらに集約（atomic add回避）"""
    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()
    
    # Stage 1
    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)
    
    sum_kernel_stage1[(n_blocks_1,)](
        input_tensor,
        partial_sums_1,
        n_elements,
        BLOCK_SIZE=BLOCK_SIZE_1,
    )
    
    # Stage 2: もう一度Tritonで集約
    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,
        )
        # 最後だけPyTorchでfloat64
        result = partial_sums_2.to(torch.float64).sum().to(torch.float32)
    else:
        # n_blocks_2が1なら直接float64で集約
        result = partial_sums_1.to(torch.float64).sum().to(torch.float32)
    
    return result

scrolls · 72 lines total

Source code from GPU Mode and the KernelBot dataset · June 9 Researcher Reciprocity License v1.0

Changes from previous submission

Against this author's previous submission submission 68217.

  ⋯ 17 unchanged lines
      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,
+ ):
+     """Stage 2: 部分和をさらに集約（atomic add回避）"""
+     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 = 1024
-     n_blocks = triton.cdiv(n_elements, BLOCK_SIZE)
+     # Stage 1
+     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)
      
-     partial_sums = torch.empty(n_blocks, device='cuda', dtype=torch.float32)
-     
-     sum_kernel_stage1[(n_blocks,)](
+     sum_kernel_stage1[(n_blocks_1,)](
          input_tensor,
-         partial_sums,
+         partial_sums_1,
          n_elements,
-         BLOCK_SIZE=BLOCK_SIZE,
+         BLOCK_SIZE=BLOCK_SIZE_1,
      )
      
-     # float64で精度を保つ
-     result = partial_sums.to(torch.float64).sum().to(torch.float32)
+     # Stage 2: もう一度Tritonで集約
+     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,
+         )
+         # 最後だけPyTorchでfloat64
+         result = partial_sums_2.to(torch.float64).sum().to(torch.float32)
+     else:
+         # n_blocks_2が1なら直接float64で集約
+         result = partial_sums_1.to(torch.float64).sum().to(torch.float32)
+     
      return result
   No newline at end of file

scrolls · 66 diff lines total

Best evidence level for this revision: reported

JSON