submission 545300

  ⋯ 5 unchanged lines
  import triton.language as tl
  from task import input_t, output_t
   
+ # Per-block private histogram approach:
+ # Phase 1: Each block builds its own 256-bin histogram (low contention)
+ # Phase 2: Sum all per-block histograms
+  
  @triton.jit
- def histogram_kernel(
-     data_ptr, output_ptr, n_elements,
+ def histogram_per_block_kernel(
+     data_ptr, hist_ptr, n_elements, n_bins: tl.constexpr,
      BLOCK_SIZE: tl.constexpr,
  ):
      pid = tl.program_id(0)
      offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
      mask = offsets < n_elements
-     vals = tl.load(data_ptr + offsets, mask=mask, other=0)
-     # Convert to int32 for indexing
-     vals = vals.to(tl.int32)
-     # Atomic add to global histogram
-     for i in range(BLOCK_SIZE):
-         if pid * BLOCK_SIZE + i < n_elements:
-             bin_idx = tl.load(data_ptr + pid * BLOCK_SIZE + i).to(tl.int32)
-             tl.atomic_add(output_ptr + bin_idx, 1)
+     vals = tl.load(data_ptr + offsets, mask=mask, other=0).to(tl.int32)
+     # Each block writes to its own 256-bin region
+     base = pid * n_bins
+     tl.atomic_add(hist_ptr + base + vals, 1, mask=mask)
   
  def custom_kernel(data: input_t) -> output_t:
      data, output = data
-     # torch.bincount is already very fast, let's just use it
-     output[...] = torch.bincount(data, minlength=256)
+     n = data.numel()
+     BLOCK_SIZE = 4096
+     n_blocks = (n + BLOCK_SIZE - 1) // BLOCK_SIZE
+     # Allocate per-block histograms
+     per_block = torch.zeros(n_blocks, 256, device=data.device, dtype=torch.int32)
+     histogram_per_block_kernel[(n_blocks,)](
+         data, per_block, n, n_bins=256, BLOCK_SIZE=BLOCK_SIZE
+     )
+     # Reduce: sum across blocks
+     output[...] = per_block.sum(dim=0)
      return output