submission 676997

  ⋯ 8 unchanged lines
  #include <torch/extension.h>
  #include <cuda_runtime.h>
   
-  
  __global__ void __launch_bounds__(256, 4)
  histogram_kernel(const uint8_t* __restrict__ data, int64_t* __restrict__ output, int N) {
-     extern __shared__ int smem_hist[];
+     extern __shared__ unsigned int smem_hist[];
      const int tid = threadIdx.x;
   
-     // Zero shared memory
-     smem_hist[tid] = 0;
-  
-     // Zero global output in-kernel (block 0 zeros all 256 bins)
+     smem_hist[tid] = 0u;
      if (blockIdx.x == 0) {
          output[tid] = 0;
      }
      __syncthreads();
   
-     // Grid-stride loop with vectorized loads (16 bytes = 16 uint8 elements per load)
      const int vec_n = N >> 4;
      const uint4* data_vec = reinterpret_cast<const uint4*>(data);
-  
      int idx = blockIdx.x * blockDim.x + tid;
      const int stride = blockDim.x * gridDim.x;
   
-     for (int i = idx; i < vec_n; i += stride) {
-         uint4 val = data_vec[i];
-         const uint8_t* b = reinterpret_cast<const uint8_t*>(&val);
+     // 2x uint4 loads with __ldg for read-only cache path
+     const int vec_n_pairs = vec_n & ~1;
+     for (int i = idx * 2; i < vec_n_pairs; i += stride * 2) {
+         uint4 val0 = __ldg(&data_vec[i]);
+         uint4 val1 = __ldg(&data_vec[i + 1]);
+         const uint8_t* b0 = reinterpret_cast<const uint8_t*>(&val0);
+         const uint8_t* b1 = reinterpret_cast<const uint8_t*>(&val1);
+         #pragma unroll
+         for (int j = 0; j < 16; j++) {
+             atomicAdd(&smem_hist[b0[j]], 1u);
+         }
+         #pragma unroll
+         for (int j = 0; j < 16; j++) {
+             atomicAdd(&smem_hist[b1[j]], 1u);
+         }
+     }
   
+     // Handle odd vec element
+     if ((vec_n & 1) && idx == 0) {
+         uint4 val = __ldg(&data_vec[vec_n - 1]);
+         const uint8_t* b = reinterpret_cast<const uint8_t*>(&val);
          #pragma unroll
          for (int j = 0; j < 16; j++) {
-             atomicAdd(&smem_hist[b[j]], 1);
+             atomicAdd(&smem_hist[b[j]], 1u);
          }
      }
   
-     // Handle remaining elements
+     // Byte-level tail
      int tail_start = vec_n * 16;
      for (int i = tail_start + idx; i < N; i += stride) {
-         atomicAdd(&smem_hist[data[i]], 1);
+         atomicAdd(&smem_hist[__ldg(&data[i])], 1u);
      }
   
      __syncthreads();
   
-     // Write back to global memory
-     if (smem_hist[tid] > 0) {
+     if (smem_hist[tid] > 0u) {
          atomicAdd(reinterpret_cast<unsigned long long*>(&output[tid]),
                    static_cast<unsigned long long>(smem_hist[tid]));
      }
  ⋯ 1 unchanged lines
   
  void histogram_inplace(torch::Tensor data, torch::Tensor output) {
      const int N = data.numel();
-     int num_blocks = min(256, max(1, (N + 256*16 - 1) / (256*16)));
-     histogram_kernel<<<num_blocks, 256, 256*sizeof(int)>>>(
+     int num_blocks = min(256, max(1, (N + 256*32 - 1) / (256*32)));
+     histogram_kernel<<<num_blocks, 256, 256*sizeof(unsigned int)>>>(
          data.data_ptr<uint8_t>(), output.data_ptr<int64_t>(), N);
  }
  """
  ⋯ 3 unchanged lines
  """
   
  _ext = load_inline(
-     name="histogram_block_private_v1",
+     name="histogram_ldg_2xu4_256blk",
      cpp_sources=cpp_src,
      cuda_sources=cuda_src,
      functions=["histogram_inplace"],