[cub] Add static and runtime-sized shared-memory histogram privatization - #10556
[cub] Add static and runtime-sized shared-memory histogram privatization#10556robobryce wants to merge 45 commits into
Conversation
Is the direct per-block merge into the output new? In the autoresearch, did we end up using a cooperative kernel for privatized smem? Did we validate that at some point? |
You need to port all of the winning smem privatized optimizations from the raw autoresearch branch. Have you done that? No half measures. Keep it minimal but complete. |
|
Please include a perf chart in this PR using existing data (no new runs). You can make this non-draft. |
|
I ported the remaining winning SMEM-privatized changes in 9dbe9c3 and expanded the PR description with the exact limits, the low-bin/wide-counter tuning, and reused B200 performance data (including a chart). I also marked the PR ready for review. On the direct-merge question: The focused histogram and environment targets compile successfully, and pre-commit passes. Runtime remains blocked on this host by |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughSummary by CodeRabbit
WalkthroughChangesHistogram execution now supports policy-selected dynamic shared-memory privatization, separate local and output counter types, cached range decoding, and SM100 tuning. Tests cover dynamic storage, counter widths, boundary classification, and policy limits. Dynamic histogram privatization
Suggested reviewers: Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 6
🧹 Nitpick comments (5)
cub/test/catch2_test_device_histogram.cu (1)
576-583: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winsuggestion: Make this test require the dynamic kernel.
This test does not constrain the architecture or verify the selected policy. Since
cub/test/catch2_test_device_histogram_env.cuexplicitly expectsdynamic_smem_bytes == 0on SM90, it can pass through the global-memory fallback without covering runtime-sized shared-memory privatization. Gate it to SM100/B200 or assert a nonzero dynamic-memory budget before callingtest_even_and_range.As per path instructions, this test should verify the new CUB dispatch path, not only end-to-end histogram results.
Source: Path instructions
cub/cub/device/dispatch/tuning/tuning_histogram.cuh (1)
376-400: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuesuggestion:
sample_sizeandsample_size_bytesare used interchangeably in the same branch.Line 377 gates on
sample_size, Lines 385-386 switch onsample_size_bytes. They coincide forpolicy_selector_from_types, but a caller that sets them independently gets an inconsistent tier. Pick one field for the whole block.cub/cub/device/dispatch/kernels/kernel_histogram.cuh (2)
1118-1141: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuesuggestion: make the
CounterTalignment of the dynamic shared block explicit.
extern __shared__ unsigned char dynamic_smem[]reinterpreted asCounterT*relies on the dynamic-smem base alignment. Declaring the extern array asCounterT(oralignas(alignof(CounterT))) documents the requirement and survives a future 8-byte counter type.
989-999: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valuesuggestion: Hoist this launch-bounds expression into a
constexprhelper for readability; the0minBlocksPerMultiprocessorform is valid here, so the nested conditional is the only real issue.cub/cub/agent/agent_histogram.cuh (1)
709-715: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valuesuggestion: every thread redundantly writes the same pointer table into shared memory.
temp_storage.histograms[ch] = pis executed identically by all threads. It's correct (each thread reads back only values it wrote itself), but it burnsNumActiveChannelsshared stores per thread and shared space for data that is a pure function ofdyn_smem_histogram_baseandnum_privatized_bins. Keeping the per-channel bases in registers (a small localCounterT* [NumActiveChannels]view) removes both.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 61cfd019-a863-4ce6-9092-ba7e9011edf7
📒 Files selected for processing (6)
cub/cub/agent/agent_histogram.cuhcub/cub/device/dispatch/dispatch_histogram.cuhcub/cub/device/dispatch/kernels/kernel_histogram.cuhcub/cub/device/dispatch/tuning/tuning_histogram.cuhcub/test/catch2_test_device_histogram.cucub/test/catch2_test_device_histogram_env.cu
|
You gave me a graph, which is nice, but that's not the one I had in mind. I want a graph of the speedup relative to baseline. Use the existing data. |
Whatever this is fix it you have root my dude. |
There was a problem hiding this comment.
Pull request overview
This PR extends CUB’s DeviceHistogram implementation to support a runtime-sized, shared-memory–privatized histogram path (backed by extern __shared__) for bin counts beyond the existing compile-time 256-bin SMEM tier, with SM100/B200-specific tuning for dynamic shared-memory budgeting and related policy selection updates.
Changes:
- Add a new histogram sweep kernel variant that places privatized counters in dynamic shared memory and merges per-block counters directly into the final output.
- Extend histogram tuning/policy plumbing to carry a per-architecture dynamic-SMEM budget (plus static-tier launch-shape overrides) and select the dynamic-SMEM path when within tuned limits.
- Add/extend tests to cover policy properties and include correctness coverage for larger-bin configurations (including wide counters).
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
cub/test/catch2_test_device_histogram.cu |
Adds a new test case intended to exercise correctness for larger bin counts (and wide counters) relevant to the dynamic-SMEM path. |
cub/test/catch2_test_device_histogram_env.cu |
Updates policy serialization/property tests for new HistogramPolicy fields and adds checks for SM100 dynamic-SMEM budget propagation (including legacy selector path). |
cub/cub/device/dispatch/tuning/tuning_histogram.cuh |
Extends HistogramPolicy with dynamic-SMEM and static-tier shape fields; adds SM100 dynamic-SMEM budget constant and tuning updates. |
cub/cub/device/dispatch/kernels/kernel_histogram.cuh |
Introduces decode-op precompute and MRU bracket-cache support for RANGE classification; adds a dynamic-SMEM sweep kernel and adjusts static-tier launch bounds/policy usage. |
cub/cub/device/dispatch/dispatch_histogram.cuh |
Adds a dynamic shared-memory kernel tag and selection logic; computes dynamic-SMEM launch bytes, sets max dynamic-SMEM attribute, and removes unused global privatization allocation when using dynamic SMEM. |
cub/cub/agent/agent_histogram.cuh |
Adds UseDynamicSmemHistogram mode for AgentHistogram to accept privatized histogram storage from dynamic shared memory and enables bracket-cache-aware bin selection in that mode. |
You can drop the discussion of why there's no cooperative combine kernel. |
9dbe9c3 to
f223e6c
Compare
5ba1d53 to
47e8690
Compare
|
CI root cause and fix: The SM100 policy was tuned on B200 and can request more opt-in dynamic shared memory than an RTX PRO 6000 kernel can actually receive. Host dispatch previously applied the policy limit directly with Host dispatch now queries the instantiated kernel's actual dynamic-SMEM capacity, requires the runtime allocation to fit both that hardware limit and the policy limit, and opts in to the smaller of the two limits. Requests that do not fit use the existing GMEM path. Device dispatch retains its existing capacity query. I also marked the device-only Fresh validation passed the eight applicable SM100 CUB histogram/environment tests on B200 and |
|
/ok to test 47e8690 |
Why not just introduce a separate tuning policy for SM120 instead? You can have the fallback that checks smem capacity as a fallback as well. |
🔬 CUB benchmark SASS comparisonHow to request a benchmark run
Targets with a SASS change
|
|
Agreed. I added an explicit SM120 dynamic-SMEM budget in SM100 retains its 227 KiB policy budget. SM120 now uses its 99 KiB per-block budget: 99 KiB for single-channel, 64 KiB for two-channel EVEN, and 96 KiB for three- and four-channel EVEN. Multi-channel RANGE remains below the SM120 ceiling already. The runtime kernel-capacity check remains as the defensive fallback for kernel-specific limits and future devices. I added compile-time policy checks for SM120 and reran all three histogram environment test binaries; all passed (785/40, 236/19, and 631/28 assertions/test cases). |
🥳 CI Workflow Results🟩 Finished in 1h 16m: Pass: 100%/288 | Total: 3d 05h | Max: 49m 29s | Hits: 85%/252120See results here. |
Why
CUB's existing shared-memory privatized histogram path uses compile-time-sized storage and falls back to per-block global-memory privatization when the histogram exceeds that static capacity. On newer GPUs, many larger histograms still fit in opt-in shared memory.
This PR keeps the low-bin static kernel and adds a runtime-sized shared-memory kernel for that larger on-chip range. RANGE classification is split so the static kernel remains lean while the runtime-sized kernel can use interpolation and a per-thread bracket cache.
What changed
extern __shared__storage and include it in occupancy selection.SearchTransformfor static-SMEM RANGE and for GMEM RANGE below the cached-search policy threshold.CachedSearchTransformfor runtime-sized SMEM RANGE and policy-selected block-private GMEM RANGE kernels.Dispatch overview
flowchart TD A[Histogram request] --> B[Compute one channel's private-histogram footprint:<br/>bins per channel x sizeof local counter] B --> C{Does that byte count fit the selected policy's<br/>compile-time static-SMEM budget?} C -- Yes --> D[Static shared-memory mode<br/>compile-time 1,024-byte private array per active channel] D --> D1{Histogram kind} D1 -- EVEN --> D2[ScaleTransform] D1 -- RANGE --> D3[SearchTransform<br/>UpperBound binary search] C -- No --> E[Compute the complete runtime-sized allocation:<br/>bins per channel x active channels x sizeof local counter] E --> F{Does that byte count fit both the selected policy's<br/>dynamic-SMEM budget and this GPU/kernel's capacity?} F -- Yes --> G[Dynamic shared-memory mode<br/>allocate exactly the requested private counters] G --> G1{Histogram kind} G1 -- EVEN --> G2[ScaleTransform] G1 -- RANGE --> G3[CachedSearchTransform<br/>piecewise-linear interpolation<br/>plus per-thread MRU bracket] F -- No --> H[Existing global-memory privatization fallback] H --> H1{Histogram kind} H1 -- EVEN --> H2[ScaleTransform] H1 -- RANGE --> H3{Cached-search policy threshold met?} H3 -- No --> H4[SearchTransform<br/>UpperBound binary search] H3 -- Yes --> H5[CachedSearchTransform<br/>piecewise-linear interpolation<br/>plus per-thread MRU bracket]The host selector reasons entirely in bytes. It first compares the per-channel private histogram against the 1,024-byte compile-time static allocation. This means 4-byte counters use static SMEM through 256 bins/channel, while 8-byte counters use it through 128 bins/channel. Counter width is not otherwise special-cased.
If the static allocation does not fit, dispatch computes the exact runtime-sized per-block requirement:
If that allocation fits both the selected policy budget and the instantiated kernel's dynamic-SMEM capacity on the current GPU, dispatch launches the runtime-sized kernel. Otherwise it uses the existing GMEM-private fallback. This keeps architecture tuning in the policy while allowing the same architecture policy to run correctly on products with different opt-in shared-memory capacities.
On SM100, the dynamic single-channel budget is 228,352 bytes after reserving 4,096 bytes for other static kernel state. Multi-channel RANGE uses a measured 8,192-byte budget per active channel, or 2,048 bins/channel with 32-bit local counters. Multi-channel EVEN uses a measured 32,768-byte budget per active channel, or 8,192 bins/channel with 32-bit local counters. SM120 has its own explicit 99 KiB policy budget: 99 KiB for single-channel, 64 KiB for two-channel EVEN, and 96 KiB for three- and four-channel EVEN. The runtime kernel-capacity check remains a defensive fallback for kernel-specific limits and future devices. The controlled same-revision experiment below is the evidence for selecting the dynamic-SMEM tier.
The compile-time static allocation is 1,024 bytes per active channel for every default policy. The earlier reviewed version raised most SM100 configurations to 2,048 bytes, which changed the generated kernel even for 8–256 runtime bins and incorrectly kept 512-bin, 32-bit-counter requests in the static tier. Restoring the established 1,024-byte allocation removes that mechanism. At 512 bins, every eligible 32-bit-counter configuration, including multi-channel EVEN, enters the runtime-sized tier.
Policy selection
flowchart TD A[Sample and counter types, channel shape, histogram kind] --> B[policy_selector_from_types] B --> C{Selected compute capability} C -->|SM100+| D[Construct complete SM100 HistogramPolicy] C -->|SM90-SM99| E[Construct complete SM90 HistogramPolicy] C -->|Pre-SM90| F[Construct complete fallback HistogramPolicy] D --> G{Privatization mode} E --> G F --> G G -->|GMEM| H[Use gmem HistogramPrivatizationPolicy] G -->|static SMEM| I[Use static_smem HistogramPrivatizationPolicy] G -->|dynamic SMEM| J[Use dynamic_smem HistogramPrivatizationPolicy] H --> K[Histogram sweep kernel and AgentHistogram] I --> K J --> K L[Deprecated custom PolicyHub] --> M[policy_selector_from_hub adapter] M --> GHost dispatch, launch bounds, and
AgentHistogramselect the same member of the same constexprHistogramPolicy; the compatibility adapter is used only for explicitly supplied deprecated custom policy hubs.Device-side child launches are different because
cudaFuncSetAttributeis host-only. Device dispatch queries both the architecture's default per-block shared-memory limit and the instantiated kernel's static shared-memory footprint, then uses dynamic SMEM only when the complete request fits the remaining capacity; otherwise it dispatches the existing GMEM specialization. For the default direct-load B200 kernel, the static footprint is zero and the 49,152-byte capacity corresponds to 6,144 bins/channel for two active channels, 4,096 for three, and 3,072 for four with 32-bit counters. Custom policies with nonzero static SMEM automatically receive a smaller dynamic limit. Host and graph launches query the instantiated kernel's actual opt-in capacity, clamp the policy limit to that capacity, and fall back to GMEM when the requested runtime allocation does not fit.Performance
The final production
.basesweep was run on a B200 with CUDA 13.3 and GCC 13.3. It compares PR commit10722be5354fff4a075a846774bacb68ba71b1b0with upstream-main commit915456f262b57793fbf5d3ccb2655ce197d9e0b3; both binaries use their unforced production policy selectors.The matrix contains 6,480 matched PR/main cells: all four single- and multi-channel EVEN/RANGE APIs, I32 and F64 samples, 1M/16M/64M/256M/1G/2B elements, 8–2,048 bins/channel, and all 15 input shapes from the autoresearch benchmark suite.
There are remaining regressions. They are concentrated in F64: the 8–256-bin static tier for single-channel RANGE and multi-channel EVEN, plus individual input shapes in the 512–2,048-bin dynamic tier for multi-channel RANGE. Sixty-eight cells are below 0.90x. I32 is substantially more robust: both RANGE APIs have no cell below 0.95x, while the EVEN APIs have 17 such cells in total.
The plots below show the geomean over all 15 input shapes at each element count and bin count. The complete output also contains one plot per input shape.
Comprehensive PR-versus-main graphs
Single-channel EVEN
Single-channel RANGE
Three-active-channel EVEN
Three-active-channel RANGE
CachedSearchTransform for block-private GMEM RANGE
A controlled same-revision experiment changes only the block-private GMEM RANGE decode operation from
SearchTransformtoCachedSearchTransform; launch policy, load policy, and storage mode are unchanged. It covers I32/F64, 16M/64M/256M elements, all 15 input shapes, and the bin counts that are in the GMEM tier.The implementation uses
CachedSearchTransformthroughout the single-channel GMEM tier for the measured I32/F64 configurations. Multi-channel I32 is robust at 1.596x geomean with a 0.952x minimum. Multi-channel F64 has regressions at 4,096 and 8,192 bins/channel, but is robust from 16,384 bins/channel upward: 89 of 90 cells win at 16,384 with a 0.973x minimum, and all 90 cells win at 32,768. The SM100 policy therefore switches multi-channel block-private GMEM RANGE toCachedSearchTransformat 16,384 bins/channel; lower bin counts retainSearchTransform. The crossover is stored in the complete histogram policy and is enabled only for the measured primitive 32/64-bit sample and 32-bit counter configurations.A controlled same-revision B200 benchmark also compared dynamic SMEM directly with block-private GMEM while holding the launch and load policy constant. It covers 120 matched states across single- and three-active-channel EVEN, I32 and F64 samples, 16M–256M pixels, 512–8,192 bins/channel, and two input entropies.
Every measured cell favored dynamic SMEM.
The final unforced selector check used 16M I32 pixels, three active channels, uniform input, and ten samples per point:
The boundary check measured 1.560 ms at 8,192 bins and 3.869 ms at 8,193 bins, confirming that the production host selector uses dynamic SMEM through 8,192 bins/channel and GMEM above it.
Validation
The branch is based directly on
upstream/mainb7aaea69a2b07e50f09e67f2962da0243e0b7c5d; the current head ise59377ebe2.Fresh CUDA 13.3/GCC 13.3 validation for the final selector passed all host, device-side, and graph variants of:
cub.test.device.histogram.lid_0,.lid_1, and.lid_2;cub.test.device.histogram_env.lid_0,.lid_1, and.lid_2;cub.test.device.histogram_env_api.lid_0.The dynamic-SMEM correctness test covers 1,024, 4,096, and 8,192 bins for both single-channel and three-active-channel APIs. The environment suite exercises the exact three-active-channel device-launch boundary at 4,096 bins, while the 8,192-bin device case exercises the GMEM fallback. Compile-time policy tests cover the SM100 and SM120 single-channel, RANGE, and two-, three-, and four-active-channel EVEN budgets, plus the cached-GMEM enable/disable and 16,384-bin crossover boundaries for I32 and F64.
The C Parallel histogram suite passed on an RTX PRO 6000. The fresh SM100 CUB build compiled the normal, host-launcher, device-launcher, graph-launcher, environment API, and legacy custom-policy targets. Eight applicable histogram, environment, and environment-API tests passed on B200. Targeted pre-commit and
git diff --checkpass on the final head.Scope
This PR contains the static and runtime-sized shared-memory privatized algorithms plus the RANGE classification and tuning required by those kernels. It excludes the high-bin shared-memory cache/direct-atomic algorithm and the separate benchmark infrastructure work.