submission 506537

  ⋯ 1 unchanged lines
  Histogram v2 — Direct global atomicAdd, block0 zeros output at start
  @MemoryCoalesced
   
- v18 bottleneck: scratch merge writes+reads 131KB extra per call.
- For 1.3MB input that's 10% overhead → fixed 1.5µs cost.
-  
- Fix: eliminate scratch entirely.
-   - Block 0 zeros output[0..255] immediately (256 stores, ~0.5µs)
-   - Block 0 sets ready flag
-   - All other blocks spin on ready flag at START (near-zero wait)
-   - All blocks accumulate into smem
-   - All blocks atomicAdd smem → output (2KB, 256 bins × num_blocks atomics)
-  
- Spin at START vs END:
-   - END spin: wait for slowest block to finish processing its data chunk → 10-15µs
-   - START spin: wait for block 0 to do 256 stores → ~0.1µs
-   ∴ START spin is ~100x cheaper
-  
- Global atomicAdd contention at merge:
-   - 128 blocks × 256 bins = 32K atomics, but spread across 256 independent addresses
-   - 128 serialized atomics per bin × ~5ns = ~640ns = 0.6µs
-   - Still cheaper than 131KB scratch write+read (~1.5µs)
+ Dynamic SM count: num_blocks = sm_count * 2, works optimally on any GPU.
+ L4: 58 SMs → 116 blocks, B200: 132 SMs → 264 blocks.
  """
   
  import torch
  ⋯ 13 unchanged lines
   
  #define BLOCK_THREADS 256
  #define NUM_BINS      256
- #define NUM_BLOCKS    116
   
  __global__ void __launch_bounds__(BLOCK_THREADS, 4)
  histogram_kernel(
      const uint4*        __restrict__ data16,
      unsigned long long* __restrict__ output,
-     volatile int*       __restrict__ flag,  // [1], reset to 0 by block 0 at end
+     volatile int*       __restrict__ flag,
      int n16,
-     int n)
+     int n,
+     int num_blocks)
  {
-     // Block 0: zero output immediately, then signal ready
-     // All others: spin until ready, then start work
      if (blockIdx.x == 0) {
          output[threadIdx.x] = 0ULL;
          __threadfence();
  ⋯ 10 unchanged lines
      __syncthreads();
   
      const int tid    = blockIdx.x * BLOCK_THREADS + threadIdx.x;
-     const int stride = NUM_BLOCKS * BLOCK_THREADS;
+     const int stride = num_blocks * BLOCK_THREADS;
   
      for (int i = tid; i < n16; i += stride) {
          uint4 v = __ldg(data16 + i);
  ⋯ 20 unchanged lines
   
      __syncthreads();
   
-     // Merge smem → global (output already zeroed)
      atomicAdd(&output[threadIdx.x], (unsigned long long)smem[threadIdx.x]);
   
-     // Block 0 resets flag for next call (it's guaranteed last to finish
-     // since it did the most work: zeroing + full data processing)
      if (blockIdx.x == 0 && threadIdx.x == 0)
          *flag = 0;
  }
   
+ static int NUM_BLOCKS_CACHED = 0;
+  
  void histogram_cuda(
      torch::Tensor data,
      torch::Tensor output,
      torch::Tensor flag)
  {
+     if (NUM_BLOCKS_CACHED == 0) {
+         int dev, sm_count;
+         cudaGetDevice(&dev);
+         cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev);
+         NUM_BLOCKS_CACHED = sm_count * 2;
+     }
      const int n   = data.numel();
      const int n16 = n / 16;
-     histogram_kernel<<<NUM_BLOCKS, BLOCK_THREADS>>>(
+     histogram_kernel<<<NUM_BLOCKS_CACHED, BLOCK_THREADS>>>(
          (const uint4*)data.data_ptr<uint8_t>(),
          (unsigned long long*)output.data_ptr<int64_t>(),
          (volatile int*)flag.data_ptr<int>(),
-         n16, n
+         n16, n, NUM_BLOCKS_CACHED
      );
  }
  """
   
  module = load_inline(
-     name='histogram_v19',
+     name='histogram_v21',
      cpp_sources=cpp_source,
      cuda_sources=cuda_source,
      functions=['histogram_cuda'],
  ⋯ 1 unchanged lines
      extra_cuda_cflags=['-O3', '--use_fast_math', '-std=c++17'],
  )
   
- # Single int flag, reset to 0 by kernel after each call
  _flag = torch.zeros(1, device='cuda', dtype=torch.int32)
   
  def custom_kernel(data: tuple) -> torch.Tensor: