Skip to content
KernelIndex
Search⌘K
Sign in

submission 68239

koshibat · python · License unknown
Kernel source · 70 lines ↓

Use it

Vendorable · source mirrored · license unknownView source →

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

memory.py
curl "https://kernelindex.com/api/v1/implementations/kernelbot-vectorsum-v2-68239?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
149.7µs
#44 of 96
2025-11-08
Run →

Reported · How evidence levels are derived →

Source and license

sourceavailable
revision digestsha256:45a2121b21d280a5a58bac4eb47d097f78e9e618529aba1a81d25e1efa2667a5
license declaredunknown
license concludedunknown
authorskoshibat
imported2026-08-15

Techniques

Extracted from the mirrored source by pattern, never inferred. Each row cites its line.

num-warps = 8num_warps=8,

Kernel source

memory.py70 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, eviction_policy='evict_first')
    block_sum = tl.sum(data)
    tl.store(partial_sums_ptr + pid, block_sum, eviction_policy='evict_first')

@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, eviction_policy='evict_first')
    block_sum = tl.sum(data)
    tl.store(partial_sums2_ptr + pid, block_sum, eviction_policy='evict_first')

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)
    
    sum_kernel_stage1[(n_blocks_1,)](
        input_tensor,
        partial_sums_1,
        n_elements,
        BLOCK_SIZE=BLOCK_SIZE_1,
        num_warps=8,
    )
    
    BLOCK_SIZE_2 = 2048
    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=8,
        )
        result = partial_sums_2.to(torch.float64).sum().to(torch.float32)
    else:
        result = partial_sums_1.to(torch.float64).sum().to(torch.float32)
    
    return result
scrolls · 70 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 68223.

  ⋯ 13 unchanged lines
      offsets = block_start + tl.arange(0, BLOCK_SIZE)
      mask = offsets < n_elements
      
-     data = tl.load(input_ptr + offsets, mask=mask, other=0.0)
+     # キャッシュヒントを追加
+     data = tl.load(input_ptr + offsets, mask=mask, other=0.0, eviction_policy='evict_first')
      block_sum = tl.sum(data)
-     tl.store(partial_sums_ptr + pid, block_sum)
+     tl.store(partial_sums_ptr + pid, block_sum, eviction_policy='evict_first')
   
  @triton.jit
  def sum_kernel_stage2(
  ⋯ 7 unchanged lines
      offsets = block_start + tl.arange(0, BLOCK_SIZE)
      mask = offsets < n_partials
      
-     data = tl.load(partial_sums_ptr + offsets, mask=mask, other=0.0)
+     data = tl.load(partial_sums_ptr + offsets, mask=mask, other=0.0, eviction_policy='evict_first')
      block_sum = tl.sum(data)
-     tl.store(partial_sums2_ptr + pid, block_sum)
+     tl.store(partial_sums2_ptr + pid, block_sum, eviction_policy='evict_first')
   
  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,
  ⋯ 2 unchanged lines
          num_warps=8,
      )
      
-     # Stage 2
      BLOCK_SIZE_2 = 2048
      n_blocks_2 = triton.cdiv(n_blocks_1, BLOCK_SIZE_2)
scrolls · 44 diff lines total

Best evidence level for this revision: reported

JSON