submission 511431

- # kernel.py
- """
- Histogram (uint8 -> int64[256]) implemented in Triton.
-  
- Pipeline / fusion note:
- - This operation is inherently a global reduction into 256 bins using atomic adds.
- - We additionally must *overwrite* the output (set bins to 0) before accumulation.
- - These two stages cannot be safely fused into a single kernel because Triton provides no
-   global barrier across programs; clearing + accumulating in one launch would race.
- So we use two Triton kernels:
-   (1) zero-out the 256-bin output
-   (2) atomic accumulation over the input
- """
-  
- from __future__ import annotations
-  
  import torch
  import triton
  import triton.language as tl
   
   
  @triton.jit
- def _zero_256_i64_kernel(out_ptr,  # *i64
-                          BLOCK: tl.constexpr):
-     pid = tl.program_id(0)
-     # Only one program is expected, but keep pid in case of accidental larger grids.
-     offs = pid * BLOCK + tl.arange(0, BLOCK)
-     mask = offs < 256
-     tl.store(out_ptr + offs, tl.zeros([BLOCK], dtype=tl.int64), mask=mask)
+ def _histogram_kernel(data_ptr, out_ptr, SIZE,
+                       BLOCK_SIZE: tl.constexpr,
+                       SUB_BLOCK: tl.constexpr):
+     """
+     Compute a 256-bin histogram for uint8 input values [0..255].
   
+     Each Triton program ("block") processes a contiguous chunk of the input.
+     Within a block, we build a local histogram in SRAM/registers by iterating
+     over the chunk in SUB_BLOCK slices to keep working set small. We use a
+     broadcasted equality (values[:, None] == bins[None, :]) masked by valid
+     positions to count occurrences per bin, then reduce along the values axis.
+     Finally, we atomically accumulate the per-block histogram into the global
+     output buffer to combine results across blocks.
   
- @triton.jit
- def _hist_u8_to_i64_kernel(x_ptr,  # *u8
-                            out_ptr,  # *i64
-                            n_elements: tl.int32,
-                            BLOCK_SIZE: tl.constexpr):
-     pid = tl.program_id(0)
-     offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
-     mask = offs < n_elements
+     Notes on fusion:
+     - Counting and global accumulation are fused in a single kernel to avoid
+       intermediate buffers and multiple launches. The only synchronization is
+       the final atomic adds per bin per block, minimizing contention compared
+       to per-element atomics.
+     """
+     # Program id and block range
+     pid = tl.program_id(axis=0)
+     block_start = pid * BLOCK_SIZE
   
-     # Load uint8 values; cast to int32 for indexing.
-     x = tl.load(x_ptr + offs, mask=mask, other=0).to(tl.int32)  # in [0, 255]
+     # Local 256-bin histogram in int32 (sufficient for counts up to SIZE per block)
+     local_hist = tl.zeros((256,), dtype=tl.int32)
   
-     # Atomic add into global histogram bins (int64 output).
-     ones = tl.full([BLOCK_SIZE], 1, dtype=tl.int64)
-     tl.atomic_add(out_ptr + x, ones, mask=mask)
+     # Bin indices [0..255] for broadcast comparisons
+     bins = tl.arange(0, 256)
   
+     # Process the chunk in SUB_BLOCK slices
+     for start in tl.range(0, BLOCK_SIZE, SUB_BLOCK, num_stages=2):
+         offs = block_start + start + tl.arange(0, SUB_BLOCK)
+         mask = offs < SIZE
+         # Load uint8 values; out-of-bounds masked to 0 and excluded by mask in the equality
+         vals = tl.load(data_ptr + offs, mask=mask, other=tl.zeros((), dtype=tl.uint8))
+         # Broadcast compare: shape [SUB_BLOCK, 256], mask invalid rows
+         eq = (vals[:, None] == bins[None, :]) & mask[:, None]
+         # Sum over rows (values) to produce counts per bin; cast bool -> int32 before sum
+         local_hist += tl.sum(eq.to(tl.int32), 0)
   
- def kernel_function(data, out: torch.Tensor | None = None) -> torch.Tensor:
-     """
-     Compute histogram of a CUDA uint8 1D tensor into 256 bins (int64), matching:
-         out[...] = torch.bincount(data, minlength=256)
+     # Atomically add local histogram to global output (int64 bins)
+     out_ptrs = out_ptr + bins
+     tl.atomic_add(out_ptrs, local_hist.to(tl.int64))
   
-     Allowed wrapper work: argument checks, allocation, and Triton launches only.
+  
+ def kernel_function(data: torch.Tensor, output: torch.Tensor = None):
      """
-     # Support alternate calling convention kernel_function((data, out))
-     if out is None and isinstance(data, (tuple, list)) and len(data) == 2:
-         data, out = data
+     Compute a 256-bin histogram over uint8 values [0..255] using a Triton kernel.
   
+     Fused stages:
+     - Single-pass counting within each block (local per-block histogram)
+     - Atomic accumulation into the global 256-bin output
+     This avoids multiple kernels or intermediate reduction buffers.
+  
+     Args:
+         data: 1-D tensor of dtype torch.uint8 on CUDA device.
+         output: Optional preallocated 1-D tensor of length 256, dtype torch.int64 on the same device.
+                 If provided, it will be zeroed and written in-place. If not provided, one is allocated.
+  
+     Returns:
+         1-D torch.Tensor of shape [256], dtype torch.int64 on the same device as input.
+  
+     Runtime constraints adhered:
+     - Wrapper performs only validation/allocation/launch; all computation is in Triton.
+     - No torch.nn, torch.nn.functional, or PyTorch compute ops are used to form the histogram.
+     """
+     # Basic validation
      if not isinstance(data, torch.Tensor):
-         raise TypeError(f"data must be a torch.Tensor, got {type(data)}")
-     if not data.is_cuda:
-         raise ValueError("data must be a CUDA tensor")
+         raise TypeError("data must be a torch.Tensor")
+     if data.device.type != "cuda":
+         raise RuntimeError("CUDA device required")
      if data.dtype != torch.uint8:
-         raise ValueError(f"data must be torch.uint8, got {data.dtype}")
+         raise TypeError(f"data dtype must be torch.uint8, got {data.dtype}")
      if data.dim() != 1:
-         raise ValueError(f"data must be 1D, got shape={tuple(data.shape)}")
-     if not data.is_contiguous():
-         data = data.contiguous()
+         raise ValueError(f"data must be 1-D, got shape {tuple(data.shape)}")
   
-     if out is None:
-         out = torch.empty((256,), device=data.device, dtype=torch.int64)
-     if not isinstance(out, torch.Tensor):
-         raise TypeError(f"out must be a torch.Tensor, got {type(out)}")
-     if out.device != data.device:
-         raise ValueError("out must be on the same device as data")
-     if out.dtype != torch.int64:
-         raise ValueError(f"out must be torch.int64, got {out.dtype}")
-     if out.shape != (256,):
-         raise ValueError(f"out must have shape (256,), got {tuple(out.shape)}")
-     if not out.is_contiguous():
-         raise ValueError("out must be contiguous")
+     SIZE = data.numel()
   
-     # Stage 1: overwrite output bins with zeros (in Triton, not PyTorch).
-     _zero_256_i64_kernel[(1,)](out, BLOCK=256, num_warps=1)
+     # Prepare output buffer
+     if output is None:
+         output = torch.zeros(256, device=data.device, dtype=torch.int64)
+     else:
+         if output.device != data.device:
+             raise RuntimeError("output must be on the same device as data")
+         if output.dtype != torch.int64 or output.numel() != 256 or output.dim() != 1:
+             raise ValueError("output must be a 1-D tensor of length 256 and dtype torch.int64")
+         # Ensure starting from zero to avoid accumulation on garbage values
+         output.zero_()
   
-     # Stage 2: atomic histogram accumulation.
-     n_elements = data.numel()
-     if n_elements == 0:
-         return out
+     # Configure launch
+     BLOCK_SIZE = 1024  # power-of-two for good coalescing
+     SUB_BLOCK = 128    # tile slice to keep working set small
+     grid = (triton.cdiv(SIZE, BLOCK_SIZE),)
   
-     BLOCK_SIZE = 1024
-     grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
-     _hist_u8_to_i64_kernel[grid](data, out, n_elements, BLOCK_SIZE=BLOCK_SIZE, num_warps=4)
+     # Launch Triton kernel
+     _histogram_kernel[grid](
+         data, output, SIZE,
+         BLOCK_SIZE=BLOCK_SIZE,
+         SUB_BLOCK=SUB_BLOCK,
+         num_warps=4,
+         num_stages=2,
+     )
   
-     return out
+     return output
   
  import inspect