submission 569139

- import torch
  import triton
  import triton.language as tl
+ import torch
   
   
  @triton.jit
- def _histogram_kernel(data_ptr, out_ptr, SIZE,
-                       BLOCK_SIZE: tl.constexpr,
-                       SUB_BLOCK: tl.constexpr):
+ def _histogram_kernel(
+     data_ptr,
+     output_ptr,
+     n_elements,
+     BLOCK_SIZE: 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.
-  
-     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.
+     Fused histogram kernel using tl.histogram with per-block privatization.
+     
+     Stage 1: Load a block of input data, compute local histogram via tl.histogram.
+              Correct for masked-out elements that default to 0 and inflate bin 0.
+     Stage 2: Flush the corrected local histogram to global output via atomic adds.
      """
-     # Program id and block range
-     pid = tl.program_id(axis=0)
+     pid = tl.program_id(0)
+  
      block_start = pid * BLOCK_SIZE
+     offsets = block_start + tl.arange(0, BLOCK_SIZE)
+     mask = offsets < n_elements
   
-     # Local 256-bin histogram in int32 (sufficient for counts up to SIZE per block)
-     local_hist = tl.zeros((256,), dtype=tl.int32)
+     # Load input data; masked-out lanes get value 0
+     data = tl.load(data_ptr + offsets, mask=mask, other=0)
+     bin_indices = data.to(tl.int32)
   
-     # Bin indices [0..255] for broadcast comparisons
-     bins = tl.arange(0, 256)
+     # tl.histogram counts ALL BLOCK_SIZE elements, including masked-out zeros
+     local_hist = tl.histogram(bin_indices, 256)
+     local_hist = local_hist.to(tl.int64)
   
-     # 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)
+     # Count how many elements were masked out (padded with 0)
+     # These were incorrectly counted in bin 0, so subtract them
+     num_valid = tl.sum(mask.to(tl.int64))
+     num_padding = BLOCK_SIZE - num_valid
   
-     # Atomically add local histogram to global output (int64 bins)
-     out_ptrs = out_ptr + bins
-     tl.atomic_add(out_ptrs, local_hist.to(tl.int64))
+     # Correct bin 0: subtract the padding elements that were counted as 0
+     bin_offsets = tl.arange(0, 256)
+     correction = tl.where(bin_offsets == 0, num_padding, tl.zeros((256,), dtype=tl.int64))
+     local_hist = local_hist - correction
   
+     # Stage 2: Flush local histogram to global output via atomic add
+     tl.atomic_add(output_ptr + bin_offsets, local_hist, mask=local_hist > 0)
   
- def kernel_function(data: torch.Tensor, output: torch.Tensor = None):
-     """
-     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.
-  
+ def kernel_function(data: torch.Tensor, output: torch.Tensor) -> torch.Tensor:
+     """
+     Wrapper for the Triton histogram kernel.
+     
      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.
-  
+         data: Input tensor of dtype uint8 on CUDA, containing values in [0, 255].
+         output: Pre-allocated output tensor of shape (256,) and dtype int64 on CUDA.
+     
      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.
+         output tensor filled with histogram bin counts.
      """
-     # Basic validation
-     if not isinstance(data, torch.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 TypeError(f"data dtype must be torch.uint8, got {data.dtype}")
-     if data.dim() != 1:
-         raise ValueError(f"data must be 1-D, got shape {tuple(data.shape)}")
+     assert data.device.type == 'cuda', "Input must be on CUDA"
+     assert data.dtype == torch.uint8, "Input must be uint8"
+     assert output.shape == (256,), "Output must have shape (256,)"
+     assert output.dtype == torch.int64, "Output must be int64"
   
-     SIZE = data.numel()
+     n_elements = data.numel()
   
-     # 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_()
+     # Zero the output histogram bins before accumulation
+     output.zero_()
   
-     # 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 = 2048
+     grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
   
-     # Launch Triton kernel
      _histogram_kernel[grid](
-         data, output, SIZE,
+         data,
+         output,
+         n_elements,
          BLOCK_SIZE=BLOCK_SIZE,
-         SUB_BLOCK=SUB_BLOCK,
-         num_warps=4,
-         num_stages=2,
      )
   
      return output
   
  import inspect
-  
  def custom_kernel(input):
      sig = inspect.signature(kernel_function)
      num_params = len(sig.parameters)
-  
      if len(input) == num_params:
          return kernel_function(*input)
      return kernel_function(input)
   
-  
- # Ensure deterministic cuBLAS.
  import os
  if os.environ.get("CUBLAS_WORKSPACE_CONFIG", "") not in (":4096:8", ":16:8"):
      os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
-