submission 611854

  ⋯ 9 unchanged lines
  #include <cuda_runtime.h>
  #include <cstdint>
   
- // Per-warp private histograms in shared memory to minimize atomic contention.
- // Each warp gets its own 256-bin histogram. At the end, we reduce across warps
- // and atomicAdd to global memory.
-  
- __global__ void histogram_kernel(
+ __global__ __launch_bounds__(256)
+ void histogram_kernel(
      const uint8_t* __restrict__ data,
      int64_t* __restrict__ output,
      const int64_t n
  ) {
-     const int WARPS_PER_BLOCK = blockDim.x / 32;
-     extern __shared__ uint32_t smem[];
+     constexpr int WARPS = 8;
+     __shared__ uint32_t smem[WARPS * 256];
   
-     const int warp_id = threadIdx.x / 32;
-     const int lane_id = threadIdx.x & 31;
+     const int warp_id = threadIdx.x >> 5;
   
-     // Zero shared memory - each warp zeros its own histogram
-     uint32_t* my_hist = smem + warp_id * 256;
+     // Zero shared memory
      #pragma unroll
-     for (int i = lane_id; i < 256; i += 32) {
-         my_hist[i] = 0;
+     for (int i = threadIdx.x; i < WARPS * 256; i += 256) {
+         smem[i] = 0;
      }
      __syncthreads();
   
-     // Global index and stride
-     const int64_t tid = blockIdx.x * blockDim.x + threadIdx.x;
-     const int64_t grid_stride = (int64_t)blockDim.x * gridDim.x;
+     uint32_t* my_hist = smem + warp_id * 256;
   
-     // Process 16 bytes (16 uint8 values) at a time using uint4 loads
-     const int64_t n16 = n / 16;
+     const int64_t tid = blockIdx.x * 256 + threadIdx.x;
+     const int64_t grid_stride = 256LL * gridDim.x;
+  
+     const int64_t n16 = n >> 4;
      const uint4* data_vec = (const uint4*)data;
   
      for (int64_t i = tid; i < n16; i += grid_stride) {
          uint4 vals = __ldg(&data_vec[i]);
   
-         #pragma unroll
-         for (int shift = 0; shift < 32; shift += 8) {
-             atomicAdd(&my_hist[(vals.x >> shift) & 0xFF], 1u);
-             atomicAdd(&my_hist[(vals.y >> shift) & 0xFF], 1u);
-             atomicAdd(&my_hist[(vals.z >> shift) & 0xFF], 1u);
-             atomicAdd(&my_hist[(vals.w >> shift) & 0xFF], 1u);
-         }
+         atomicAdd(&my_hist[(vals.x      ) & 0xFF], 1u);
+         atomicAdd(&my_hist[(vals.x >>  8) & 0xFF], 1u);
+         atomicAdd(&my_hist[(vals.x >> 16) & 0xFF], 1u);
+         atomicAdd(&my_hist[(vals.x >> 24)       ], 1u);
+         atomicAdd(&my_hist[(vals.y      ) & 0xFF], 1u);
+         atomicAdd(&my_hist[(vals.y >>  8) & 0xFF], 1u);
+         atomicAdd(&my_hist[(vals.y >> 16) & 0xFF], 1u);
+         atomicAdd(&my_hist[(vals.y >> 24)       ], 1u);
+         atomicAdd(&my_hist[(vals.z      ) & 0xFF], 1u);
+         atomicAdd(&my_hist[(vals.z >>  8) & 0xFF], 1u);
+         atomicAdd(&my_hist[(vals.z >> 16) & 0xFF], 1u);
+         atomicAdd(&my_hist[(vals.z >> 24)       ], 1u);
+         atomicAdd(&my_hist[(vals.w      ) & 0xFF], 1u);
+         atomicAdd(&my_hist[(vals.w >>  8) & 0xFF], 1u);
+         atomicAdd(&my_hist[(vals.w >> 16) & 0xFF], 1u);
+         atomicAdd(&my_hist[(vals.w >> 24)       ], 1u);
      }
   
-     // Handle remaining elements
+     // Tail
      {
          int64_t tail_start = n16 * 16;
          for (int64_t i = tail_start + tid; i < n; i += grid_stride) {
  ⋯ 3 unchanged lines
   
      __syncthreads();
   
-     // Reduce across warps and write to global memory
-     for (int bin = threadIdx.x; bin < 256; bin += blockDim.x) {
+     if (threadIdx.x < 256) {
          uint32_t total = 0;
-         for (int w = 0; w < WARPS_PER_BLOCK; w++) {
-             total += smem[w * 256 + bin];
+         #pragma unroll
+         for (int w = 0; w < WARPS; w++) {
+             total += smem[w * 256 + threadIdx.x];
          }
          if (total > 0) {
-             atomicAdd((unsigned long long*)&output[bin], (unsigned long long)total);
+             atomicAdd((unsigned long long*)&output[threadIdx.x], (unsigned long long)total);
          }
      }
  }
   
  torch::Tensor histogram_cuda(torch::Tensor data, torch::Tensor output) {
      const int64_t n = data.numel();
-     output.zero_();
+     cudaMemsetAsync(output.data_ptr<int64_t>(), 0, 256 * sizeof(int64_t));
   
-     const int threads = 256;
-     const int max_blocks = 160 * 4;
-     const int64_t elements_per_block = threads * 16;
-     int blocks = (n + elements_per_block - 1) / elements_per_block;
-     if (blocks > max_blocks) blocks = max_blocks;
-     if (blocks < 1) blocks = 1;
+     int blocks = 320;
+     if (n < 256 * 16 * blocks) {
+         blocks = (n + 256 * 16 - 1) / (256 * 16);
+         if (blocks < 1) blocks = 1;
+     }
   
-     const int warps_per_block = threads / 32;
-     const int smem_size = warps_per_block * 256 * sizeof(uint32_t);
-  
-     histogram_kernel<<<blocks, threads, smem_size>>>(
+     histogram_kernel<<<blocks, 256>>>(
          data.data_ptr<uint8_t>(),
          output.data_ptr<int64_t>(),
          n
  ⋯ 8 unchanged lines
  """;
   
  module = load_inline(
-     name="histogram_kernel",
+     name="histogram_kernel_v15",
      cpp_sources=[cpp_source],
      cuda_sources=[cuda_source],
      functions=["histogram_cuda"],