submission 758463

  ⋯ 9 unchanged lines
   
   
  @triton.jit
- def histogram_kernel(
-     data_ptr,
-     tmp_hist_ptr,
-     N,
+ def histogram_warp_kernel(
+     data_ptr,          # *uint8
+     partial_ptr,       # *int32
+     N: tl.int32,
      BLOCK_SIZE: tl.constexpr,
      CHUNK: tl.constexpr,
+     WARP_SIZE: tl.constexpr,
  ):
+     """Per‑block, per‑warp histogram kernel.
+  
+     Each block processes a contiguous ``CHUNK`` of the input.  Within a block
+     we allocate ``NUM_WARPS`` independent 256‑bin histograms in global memory.
+     Threads update the histogram belonging to their warp using atomic adds.
      """
-     Build a per‑block 256‑bin histogram in a temporary global buffer.
-     Each block writes to its own slice of `tmp_hist` (256 int32 entries).
-     """
-     pid = tl.program_id(0)
+     pid = tl.program_id(0)                # block index
+     lane = tl.arange(0, BLOCK_SIZE)       # thread index within block
   
-     # Pointer to the slice belonging to this block (256 int32 bins)
-     tmp_ptr = tmp_hist_ptr + pid * 256
+     # Number of warps per block (BLOCK_SIZE is a multiple of 32)
+     NUM_WARPS = BLOCK_SIZE // WARP_SIZE
   
-     # Thread index inside the block
-     thread_idx = tl.arange(0, BLOCK_SIZE)
+     # Base pointer for this block's per‑warp histograms (flattened layout)
+     # Layout: [block, warp, bin]
+     block_base = partial_ptr + pid * NUM_WARPS * 256
   
-     # Base offset for this thread: each thread processes CHUNK elements spaced by BLOCK_SIZE
-     base = pid * BLOCK_SIZE * CHUNK + thread_idx
+     # -----------------------------------------------------------------
+     # Zero‑initialize each warp's histogram.
+     # Every thread writes zero to its lane bin for every warp.
+     for w in range(NUM_WARPS):
+         tl.store(block_base + w * 256 + lane, 0)
   
-     # Loop over the CHUNK elements assigned to this thread
-     for i in range(CHUNK):
-         idx = base + i * BLOCK_SIZE
-         mask = idx < N
+     # Make sure the zero stores are visible before any atomics.
+     tl.debug_barrier()
   
-         # Load a uint8 value (0‑255) and cast to int32 for indexing
-         val = tl.load(data_ptr + idx, mask=mask, other=0).to(tl.int32)
+     # -----------------------------------------------------------------
+     # Process a contiguous chunk of the input.
+     block_start = pid * CHUNK
+     ITEMS_PER_THREAD = CHUNK // BLOCK_SIZE   # compile‑time constant
   
-         # Increment the per‑block histogram atomically
-         tl.atomic_add(tmp_ptr + val, 1, mask=mask)
+     for i in range(ITEMS_PER_THREAD):
+         offset = block_start + lane + i * BLOCK_SIZE
+         mask = offset < N
+         # Load a uint8 value (masked)
+         val = tl.load(data_ptr + offset, mask=mask, other=0)
+         # Convert to a 0‑255 index
+         idx = tl.cast(val, tl.int32)
+         # Identify the warp this thread belongs to
+         warp_id = lane // WARP_SIZE
+         # Pointer to the appropriate bin in this warp's histogram
+         bin_ptr = block_base + warp_id * 256 + idx
+         # Atomically increment the bin (masked)
+         tl.atomic_add(bin_ptr, 1, mask=mask)
   
   
- @triton.jit
- def reduce_histogram_kernel(
-     tmp_hist_ptr,
-     out_ptr,
-     BLOCK_SIZE: tl.constexpr,
- ):
-     """
-     Reduce all per‑block histograms into the final output.
-     One thread per bin (0‑255) adds its bin count from this block to `out`.
-     """
-     pid = tl.program_id(0)
-  
-     # One thread per histogram bin
-     bin_idx = tl.arange(0, BLOCK_SIZE)
-     mask = bin_idx < 256
-  
-     # Load the count for this bin from this block's temporary histogram (int32 → int64)
-     cnt = tl.load(tmp_hist_ptr + pid * 256 + bin_idx, mask=mask).to(tl.int64)
-  
-     # Atomically add to the final output (int64)
-     tl.atomic_add(out_ptr + bin_idx, cnt, mask=mask)
-  
-  
  def custom_kernel(data: input_t) -> output_t:
      """
-     Compute a 256‑bin histogram of a uint8 1‑D tensor.
-     This implementation builds per‑block sub‑histograms in a temporary buffer
-     to dramatically reduce atomic contention, then reduces them into the final
-     output.
+     Compute a 256‑bin histogram of a 1‑D uint8 tensor using a Triton kernel.
+     The kernel builds per‑warp partial histograms in global memory and then
+     reduces them with a fast Torch sum.
      """
-     data_tensor, output_tensor = data
-  
-     # Ensure inputs are contiguous and on the same device
-     data_tensor = data_tensor.contiguous()
-     output_tensor = output_tensor.contiguous()
+     data_tensor, output = data
+     N = data_tensor.shape[0]
      device = data_tensor.device
   
-     # Zero the output tensor (required for atomic accumulation)
-     output_tensor.zero_()
+     # -----------------------------------------------------------------
+     # Handle empty input early.
+     if N == 0:
+         output.zero_()
+         return output
   
-     N = data_tensor.numel()
+     # -----------------------------------------------------------------
+     # Tunable launch parameters
+     BLOCK_SIZE = 256            # one thread per histogram bin
+     WARP_SIZE = 32
+     NUM_WARPS = BLOCK_SIZE // WARP_SIZE
+     # CHUNK: number of elements processed per block (must be multiple of BLOCK_SIZE)
+     CHUNK = 16384               # 64 * BLOCK_SIZE
   
-     # Kernel launch configuration
-     BLOCK_SIZE = 256          # one thread per histogram bin (256 threads)
-     CHUNK = 16                # elements processed per thread
+     # Number of blocks required to cover the input
+     num_blocks = (N + CHUNK - 1) // CHUNK
   
-     # Number of blocks needed to cover the input
-     num_blocks = max(1, triton.cdiv(N, BLOCK_SIZE * CHUNK))
-     grid = (num_blocks,)
+     # Allocate per‑warp partial histograms (int32 is sufficient for counts)
+     # Shape: [num_blocks, NUM_WARPS, 256]
+     partial = torch.empty((num_blocks, NUM_WARPS, 256), dtype=torch.int32, device=device)
   
-     # Temporary buffer for per‑block histograms (int32)
-     tmp_hist = torch.zeros((num_blocks, 256), dtype=torch.int32, device=device)
-  
-     # -----------------------------------------------------------------
-     # Phase 1: build per‑block histograms
-     histogram_kernel[grid](
+     # Launch the kernel
+     histogram_warp_kernel[(num_blocks,)](
          data_tensor,
-         tmp_hist,
+         partial,
          N,
          BLOCK_SIZE=BLOCK_SIZE,
          CHUNK=CHUNK,
-         num_warps=8,
-         num_stages=4,
+         WARP_SIZE=WARP_SIZE,
+         num_warps=NUM_WARPS,
+         num_stages=3,
      )
-     # -----------------------------------------------------------------
-     # Phase 2: reduce per‑block histograms into the final output
-     reduce_histogram_kernel[grid](
-         tmp_hist,
-         output_tensor,
-         BLOCK_SIZE=BLOCK_SIZE,
-         num_warps=8,
-         num_stages=2,
-     )
-     # -----------------------------------------------------------------
   
-     return output_tensor
+     # Reduce across blocks and warps into the final 256‑bin histogram.
+     # Cast to the output dtype (int64) during the reduction.
+     output[:] = partial.sum(dim=(0, 1), dtype=output.dtype)
+  
+     return output
  # EVOLVE-BLOCK-END