submission 759321

  ⋯ 4 unchanged lines
  import triton.language as tl
  from typing import Tuple
   
- # ------------------------------------------------------------------
- # Kernel: compute per‑block sums
- # ------------------------------------------------------------------
+ # -------------------------------------------------------------------------
+ # Tunable configuration (target Hopper H200)
+ # -------------------------------------------------------------------------
+ BLOCK_THREADS = 256               # 8 warps per block
+ VECT = 64                         # elements processed per thread
+ BLOCK_ELEMENTS = BLOCK_THREADS * VECT  # total elements per block (16384)
+  
+ # -------------------------------------------------------------------------
+ # Kernel: compute per‑block sums (masked)
+ # -------------------------------------------------------------------------
  @triton.jit
  def block_sum_kernel(
-     x_ptr,               # *float32, input vector
-     block_sum_ptr,       # *float32, per‑block sums (output)
-     N,                   # int32, total number of elements
-     BLOCK_SIZE: tl.constexpr,
+     x_ptr,               # *const float32
+     block_sums_ptr,      # *float32
+     N: tl.int32,         # total number of elements
+     BLOCK_THREADS: tl.constexpr,
+     VECT: tl.constexpr,
  ):
      pid = tl.program_id(0)                     # block index
-     block_start = pid * BLOCK_SIZE
-     offsets = block_start + tl.arange(0, BLOCK_SIZE)
-     mask = offsets < N
+     block_start = pid * BLOCK_THREADS * VECT
   
-     # Load a tile (zero‑pad out‑of‑bounds)
-     x = tl.load(x_ptr + offsets, mask=mask, other=0.0)
+     # thread and vector indices
+     tid = tl.arange(0, BLOCK_THREADS)[:, None]   # (BLOCK_THREADS, 1)
+     vec = tl.arange(0, VECT)[None, :]            # (1, VECT)
   
-     # Reduce across the block
+     offs = block_start + tid * VECT + vec
+     mask = offs < N
+  
+     # Load tile (masked)
+     x = tl.load(x_ptr + offs, mask=mask, other=0.0,
+                 eviction_policy='evict_last')
+     # Reduce across the whole tile
      block_sum = tl.sum(x)
+     # Write per‑block sum
+     tl.store(block_sums_ptr + pid, block_sum)
   
-     # Store the block sum
-     tl.store(block_sum_ptr + pid, block_sum)
   
+ # -------------------------------------------------------------------------
+ # Kernel: exclusive scan of per‑block sums (single‑block, in‑place)
+ # -------------------------------------------------------------------------
+ @triton.jit
+ def block_offsets_scan_kernel(
+     block_sums_ptr,      # *float32 (in‑place: input sums → exclusive offsets)
+     M: tl.int32,         # number of block sums (num_blocks)
+     BLOCK_THREADS: tl.constexpr,
+     VECT: tl.constexpr,
+ ):
+     tid = tl.arange(0, BLOCK_THREADS)[:, None]
+     vec = tl.arange(0, VECT)[None, :]
   
- # ------------------------------------------------------------------
- # Kernel: final inclusive scan using per‑block exclusive offsets
- # ------------------------------------------------------------------
+     offs = tid * VECT + vec
+     mask = offs < M
+  
+     # Load block sums (masked)
+     x = tl.load(block_sums_ptr + offs, mask=mask, other=0.0,
+                 eviction_policy='evict_last')
+     # Flatten to 1‑D
+     x_flat = tl.reshape(x, (BLOCK_THREADS * VECT,))
+     # Inclusive prefix sum
+     inc = tl.cumsum(x_flat)
+     # Convert to exclusive by subtracting the original value
+     excl = inc - x_flat
+     # Store exclusive offsets back (in‑place)
+     excl_2d = tl.reshape(excl, (BLOCK_THREADS, VECT))
+     tl.store(block_sums_ptr + offs, excl_2d, mask=mask)
+  
+  
+ # -------------------------------------------------------------------------
+ # Kernel: inclusive scan inside each block and add the exclusive block offset
+ # -------------------------------------------------------------------------
  @triton.jit
- def final_scan_kernel(
-     x_ptr,               # *float32, input vector
-     out_ptr,             # *float32, output vector (final result)
-     block_offsets_ptr,   # *float32, inclusive per‑block offsets
-     block_sums_ptr,      # *float32, per‑block sums
-     N,                   # int32, total number of elements
-     BLOCK_SIZE: tl.constexpr,
+ def block_scan_add_offset_kernel(
+     x_ptr,                # *const float32
+     y_ptr,                # *float32
+     offsets_ptr,          # *float32 (exclusive per‑block offsets)
+     N: tl.int32,          # total number of elements
+     BLOCK_THREADS: tl.constexpr,
+     VECT: tl.constexpr,
  ):
-     pid = tl.program_id(0)                     # block index
-     block_start = pid * BLOCK_SIZE
-     offsets = block_start + tl.arange(0, BLOCK_SIZE)
-     mask = offsets < N
+     pid = tl.program_id(0)
+     block_start = pid * BLOCK_THREADS * VECT
   
-     # Load input tile
-     x = tl.load(x_ptr + offsets, mask=mask, other=0.0)
+     tid = tl.arange(0, BLOCK_THREADS)[:, None]
+     vec = tl.arange(0, VECT)[None, :]
   
-     # Inclusive prefix sum inside the block
-     y = tl.cumsum(x)
+     offs = block_start + tid * VECT + vec
+     mask = offs < N
   
-     # Load inclusive offset for this block and its block sum
-     offset_incl = tl.load(block_offsets_ptr + pid)
-     block_sum = tl.load(block_sums_ptr + pid)
+     # Load the exclusive offset for this block (scalar)
+     offset = tl.load(offsets_ptr + pid)
   
-     # Convert inclusive offset to exclusive offset
-     offset_excl = offset_incl - block_sum
+     # Load the input tile (masked)
+     x = tl.load(x_ptr + offs, mask=mask, other=0.0,
+                 eviction_policy='evict_last')
   
-     # Apply the exclusive offset
-     y = y + offset_excl
+     # Inclusive prefix sum within the block
+     x_flat = tl.reshape(x, (BLOCK_THREADS * VECT,))
+     y_flat = tl.cumsum(x_flat) + offset
   
-     # Write results
-     tl.store(out_ptr + offsets, y, mask=mask)
+     # Store result
+     y = tl.reshape(y_flat, (BLOCK_THREADS, VECT))
+     tl.store(y_ptr + offs, y, mask=mask)
   
   
- # ------------------------------------------------------------------
- # Kernel: single‑pass scan for very small inputs (N <= BLOCK_SIZE)
- # ------------------------------------------------------------------
+ # -------------------------------------------------------------------------
+ # Fast‑path kernel for inputs that fit within a single block
+ # -------------------------------------------------------------------------
  @triton.jit
- def small_scan_kernel(
-     x_ptr,               # *float32, input vector
-     out_ptr,             # *float32, output vector
-     N,                   # int32, total number of elements
-     BLOCK_SIZE: tl.constexpr,
+ def single_block_scan_kernel(
+     x_ptr,
+     y_ptr,
+     N: tl.int32,
+     BLOCK_THREADS: tl.constexpr,
+     VECT: tl.constexpr,
  ):
-     offsets = tl.arange(0, BLOCK_SIZE)
-     mask = offsets < N
-     x = tl.load(x_ptr + offsets, mask=mask, other=0.0)
-     y = tl.cumsum(x)
-     tl.store(out_ptr + offsets, y, mask=mask)
+     # Only one block (pid == 0) is launched
+     tid = tl.arange(0, BLOCK_THREADS)[:, None]
+     vec = tl.arange(0, VECT)[None, :]
+     offs = tid * VECT + vec
+     mask = offs < N
   
+     x = tl.load(x_ptr + offs, mask=mask, other=0.0,
+                 eviction_policy='evict_last')
+     x_flat = tl.reshape(x, (BLOCK_THREADS * VECT,))
+     y_flat = tl.cumsum(x_flat)
+     y = tl.reshape(y_flat, (BLOCK_THREADS, VECT))
+     tl.store(y_ptr + offs, y, mask=mask)
   
- # ------------------------------------------------------------------
- # Public interface
- # ------------------------------------------------------------------
+  
  def custom_kernel(data: Tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
      """
-     Inclusive prefix sum (scan) over a 1‑D float32 tensor using Triton.
+     Inclusive prefix sum (scan) over a 1‑D float32 CUDA tensor using Triton.
+     The algorithm consists of three steps:
+       1) Compute per‑block sums.
+       2) Perform an exclusive scan of those block sums on‑device
+          (single‑block kernel, in‑place).
+       3) Scan each block of the original tensor and add the pre‑computed offset.
+     A fast‑path kernel handles inputs that fit within a single block.
      """
      x, output = data
   
-     # Basic sanity checks
-     assert x.is_contiguous() and output.is_contiguous(), "Both tensors must be contiguous"
+     # -----------------------------------------------------------------
+     # Sanity checks
+     # -----------------------------------------------------------------
+     assert x.is_cuda and output.is_cuda, "Triton kernels require CUDA tensors"
      assert x.dtype == torch.float32 and output.dtype == torch.float32, "Only float32 supported"
+     assert x.is_contiguous() and output.is_contiguous(), "Tensors must be contiguous"
      assert x.shape == output.shape, "Input and output must have the same shape"
   
-     N = x.shape[0]
-     if N == 0:
-         output.zero_()
+     n = x.shape[0]
+     if n == 0:
          return output
   
-     # ------------------------------------------------------------------
-     # Configuration (tuned for Hopper/H200)
-     # ------------------------------------------------------------------
-     BLOCK_SIZE = 1024                     # max threads per block
-     num_blocks = (N + BLOCK_SIZE - 1) // BLOCK_SIZE
-  
-     # ------------------------------------------------------------------
-     # Small‑input fast path
-     # ------------------------------------------------------------------
-     if N <= BLOCK_SIZE:
-         small_scan_kernel[(1,)](
-             x, output, N,
-             BLOCK_SIZE=BLOCK_SIZE,
-             num_stages=2,
+     # -----------------------------------------------------------------
+     # Fast‑path: single block (n ≤ BLOCK_ELEMENTS)
+     # -----------------------------------------------------------------
+     if n <= BLOCK_ELEMENTS:
+         grid = (1,)
+         single_block_scan_kernel[grid](
+             x,
+             output,
+             n,
+             BLOCK_THREADS=BLOCK_THREADS,
+             VECT=VECT,
+             num_warps=BLOCK_THREADS // 32,
+             num_stages=4,
          )
          return output
   
-     # ------------------------------------------------------------------
-     # Large‑input hierarchical scan
-     # ------------------------------------------------------------------
-     # 1️⃣ Compute per‑block sums
+     # -----------------------------------------------------------------
+     # General case: hierarchical scan
+     # -----------------------------------------------------------------
+     num_blocks = (n + BLOCK_ELEMENTS - 1) // BLOCK_ELEMENTS
+  
+     # Buffer for per‑block sums (will later hold exclusive offsets)
      block_sums = torch.empty(num_blocks, dtype=x.dtype, device=x.device)
+  
+     # -----------------------------------------------------------------
+     # 1️⃣ Compute per‑block sums.
+     # -----------------------------------------------------------------
      block_sum_kernel[(num_blocks,)](
-         x, block_sums, N,
-         BLOCK_SIZE=BLOCK_SIZE,
+         x,
+         block_sums,
+         n,
+         BLOCK_THREADS=BLOCK_THREADS,
+         VECT=VECT,
+         num_warps=BLOCK_THREADS // 32,
+         num_stages=4,
+     )
+  
+     # -----------------------------------------------------------------
+     # 2️⃣ Compute exclusive offsets of the block sums (in‑place).
+     # -----------------------------------------------------------------
+     block_offsets_scan_kernel[(1,)](
+         block_sums,
+         num_blocks,
+         BLOCK_THREADS=BLOCK_THREADS,
+         VECT=VECT,
+         num_warps=BLOCK_THREADS // 32,
          num_stages=2,
      )
   
-     # 2️⃣ Compute inclusive block offsets on the GPU
-     block_offsets = torch.empty_like(block_sums)
-     torch.cumsum(block_sums, dim=0, out=block_offsets)
-  
-     # 3️⃣ Final scan with exclusive offsets
-     final_scan_kernel[(num_blocks,)](
-         x, output, block_offsets, block_sums, N,
-         BLOCK_SIZE=BLOCK_SIZE,
+     # -----------------------------------------------------------------
+     # 3️⃣ Scan each block and add its exclusive offset.
+     # -----------------------------------------------------------------
+     block_scan_add_offset_kernel[(num_blocks,)](
+         x,
+         output,
+         block_sums,          # now contains exclusive offsets
+         n,
+         BLOCK_THREADS=BLOCK_THREADS,
+         VECT=VECT,
+         num_warps=BLOCK_THREADS // 32,
          num_stages=4,
      )
      return output