Handle RF split histogram shared-memory pressure - #8323
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
RAMitchell
added
improvement
There was a problem hiding this comment.
Pull request overview
This PR adds a runtime fallback in the batched-level decision tree / random forest split-histogram computation to avoid shared-memory exhaustion (or severe occupancy loss) for large n_classes * max_n_bins classification workloads by switching to a global-memory histogram while keeping split scoring logic in a single kernel.
Changes:
- Adds a
use_global_memory_histogramruntime switch tocomputeSplitKernel/launchComputeSplitKernel, selecting shared-memory vs global-memory histogram storage. - Refactors histogram/CDF conversion to operate on either shared or global histogram pointers, and updates builder logic to size/clear histogram workspace appropriately for the selected path.
- Adds an RF classification test targeting high class-count (
n_classes=80,max_n_bins=256) to exercise the fallback and ensure the fitted tree actually splits.
Reviewed changes
Copilot reviewed 12 out of 12 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| cpp/tests/sg/rf_test.cu | Adds a high-class-count RF test that exercises the global-memory histogram fallback and checks the tree splits. |
| cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh | Updates lower_bound to take a const pointer and extends launchComputeSplitKernel signature with the runtime fallback boolean. |
| cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels_impl.cuh | Implements the shared/global histogram pointer switching inside computeSplitKernel and threads the new boolean through the kernel launcher. |
| cpp/src/decisiontree/batched-levelalgo/builder.cuh | Adds the tunable heuristic threshold and selects shared vs global histogram mode at runtime, adjusting histogram workspace sizing accordingly. |
| cpp/src/decisiontree/batched-levelalgo/kernels/classification-float.cu | Updates explicit template instantiation to include use_global_memory_histogram. |
| cpp/src/decisiontree/batched-levelalgo/kernels/classification-double.cu | Updates explicit template instantiation to include use_global_memory_histogram. |
| cpp/src/decisiontree/batched-levelalgo/kernels/weighted-classification-float.cu | Updates explicit template instantiation to include use_global_memory_histogram. |
| cpp/src/decisiontree/batched-levelalgo/kernels/weighted-classification-double.cu | Updates explicit template instantiation to include use_global_memory_histogram. |
| cpp/src/decisiontree/batched-levelalgo/kernels/regression-float.cu | Updates explicit template instantiation to include use_global_memory_histogram. |
| cpp/src/decisiontree/batched-levelalgo/kernels/regression-double.cu | Updates explicit template instantiation to include use_global_memory_histogram. |
| cpp/src/decisiontree/batched-levelalgo/kernels/weighted-regression-float.cu | Updates explicit template instantiation to include use_global_memory_histogram. |
| cpp/src/decisiontree/batched-levelalgo/kernels/weighted-regression-double.cu | Updates explicit template instantiation to include use_global_memory_histogram. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughSummary by CodeRabbit
WalkthroughThis PR adds a global-memory histogram fallback for decision-tree split computation when shared memory would be exceeded. It adds path-selection logic in ChangesGlobal-memory histogram fallback
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
cpp/src/decisiontree/batched-levelalgo/builder.cuh (1)
571-614: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winHoist the histogram-path decision out of the per-column-batch loop.
shouldUseGlobalMemoryHistogram(and its inputscomputeSplitHistogramSmemSize()/computeSplitSmemSize(), plus a device-property query) don't depend oncol, yetcomputeSplitis invoked once per iteration of thefor (IdxT c = 0; c < dataset.n_sampled_cols; c += n_blks_for_cols)loop incomputeBestSplits(Line 498). The decision and both size calculations are recomputed identically on every column batch of every node-queue round. Since this decision is invariant perBuilderinstance/params, consider computing it once (e.g., in the constructor or at the top ofcomputeBestSplits) and threading the resultinguse_global_memory_histogram/smem_sizeintocomputeSplit.♻️ Sketch of hoisting the decision
- void computeSplit(IdxT col, size_t n_blocks_dimx, size_t n_large_nodes, size_t n_work_items) + void computeSplit(IdxT col, + size_t n_blocks_dimx, + size_t n_large_nodes, + size_t n_work_items, + bool use_global_memory_histogram, + size_t smem_size) { // if no instances to split, return if (n_blocks_dimx == 0) return; raft::common::nvtx::range fun_scope("Builder::computeSplit `@builder.cuh` [batched-levelalgo]"); auto n_bins = params.max_n_bins; auto n_classes = dataset.num_outputs; - auto shared_histogram_dynamic_smem_size = computeSplitHistogramSmemSize(); - auto shared_path_total_smem_size = computeSplitSmemSize(); - auto use_global_memory_histogram = shouldUseGlobalMemoryHistogram( - shared_histogram_dynamic_smem_size, shared_path_total_smem_size); - auto smem_size = use_global_memory_histogram ? computeSplitGlobalHistogramSmemSize() - : shared_histogram_dynamic_smem_size; // if columns left to be processed lesser than `n_blks_for_cols`, shrink the blocks along dimyCompute
use_global_memory_histogram/smem_sizeonce incomputeBestSplitsand pass them through.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/decisiontree/batched-levelalgo/builder.cuh` around lines 571 - 614, The histogram-path selection in computeSplit is being recomputed for every column batch even though it is invariant for a given Builder configuration. Hoist the calls to shouldUseGlobalMemoryHistogram, computeSplitHistogramSmemSize, computeSplitSmemSize, and the related device-property-dependent sizing out of computeSplit and into computeBestSplits or the Builder setup path, then pass the resulting use_global_memory_histogram and smem_size into computeSplit. Keep computeSplit focused on per-column work and use the precomputed values when building the grid and launching launchComputeSplitKernel.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@cpp/src/decisiontree/batched-levelalgo/builder.cuh`:
- Around line 571-614: The histogram-path selection in computeSplit is being
recomputed for every column batch even though it is invariant for a given
Builder configuration. Hoist the calls to shouldUseGlobalMemoryHistogram,
computeSplitHistogramSmemSize, computeSplitSmemSize, and the related
device-property-dependent sizing out of computeSplit and into computeBestSplits
or the Builder setup path, then pass the resulting use_global_memory_histogram
and smem_size into computeSplit. Keep computeSplit focused on per-column work
and use the precomputed values when building the grid and launching
launchComputeSplitKernel.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 7b0acd61-01de-4854-bd59-084c4f61b44b
📒 Files selected for processing (12)
cpp/src/decisiontree/batched-levelalgo/builder.cuhcpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuhcpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels_impl.cuhcpp/src/decisiontree/batched-levelalgo/kernels/classification-double.cucpp/src/decisiontree/batched-levelalgo/kernels/classification-float.cucpp/src/decisiontree/batched-levelalgo/kernels/regression-double.cucpp/src/decisiontree/batched-levelalgo/kernels/regression-float.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted-classification-double.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted-classification-float.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted-regression-double.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted-regression-float.cucpp/tests/sg/rf_test.cu
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
cpp/src/decisiontree/batched-levelalgo/builder.cuh (1)
589-590: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winUse
ML::checked_mul<size_t>for the histogram length product.
len_histogramsis a count-product fed directly into thecudaMemsetAsyncsize (and it bounds the buffer the kernel writes). The rest of this change (seecomputeSplitSharedMemoryConfig) already guards every count-product withML::checked_mul; this site should follow the same invariant. The leadingsize_t(n_bins)widens the arithmetic but there is no explicit overflow guard.As per path instructions: "Multiplications... of
int... whose result is passed to...cudaMalloc*... or used as asize_t/int64_tparameter... RequireML::checked_mul<size_t>(...)... or equivalent widening + explicit guard at the call site."🛡️ Proposed fix
- size_t len_histograms = size_t(n_bins) * n_classes * n_blocks_dimy * histogram_node_count; - RAFT_CUDA_TRY(cudaMemsetAsync(histograms, 0, sizeof(BinT) * len_histograms, builder_stream)); + size_t len_histograms = ML::checked_mul<std::size_t>( + std::size_t(n_bins), n_classes, n_blocks_dimy, histogram_node_count); + RAFT_CUDA_TRY(cudaMemsetAsync( + histograms, 0, ML::checked_mul<std::size_t>(sizeof(BinT), len_histograms), builder_stream));🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/decisiontree/batched-levelalgo/builder.cuh` around lines 589 - 590, The histogram buffer length in the builder logic is computed with unchecked multiplication before being passed to cudaMemsetAsync, so update the len_histograms calculation in builder.cuh to use ML::checked_mul<size_t> for the full n_bins * n_classes * n_blocks_dimy * histogram_node_count product. Keep the existing zeroing call intact, but ensure the count-product is overflow-guarded in the same style used by computeSplitSharedMemoryConfig and any other size calculations in the builder path.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@cpp/src/decisiontree/batched-levelalgo/builder.cuh`:
- Around line 589-590: The histogram buffer length in the builder logic is
computed with unchecked multiplication before being passed to cudaMemsetAsync,
so update the len_histograms calculation in builder.cuh to use
ML::checked_mul<size_t> for the full n_bins * n_classes * n_blocks_dimy *
histogram_node_count product. Keep the existing zeroing call intact, but ensure
the count-product is overflow-guarded in the same style used by
computeSplitSharedMemoryConfig and any other size calculations in the builder
path.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 331a8816-be91-42c8-9d81-bbe3e312b7c3
📒 Files selected for processing (1)
cpp/src/decisiontree/batched-levelalgo/builder.cuh
|
@RAMitchell Please avoid merging |
| handle.sync_stream_pool(); | ||
| } | ||
|
|
||
| TEST(RfTests, HighClassCountSplitHistogramFallsBackToGlobalMemory) |
There was a problem hiding this comment.
Considering that this test does not explicitly test whether the GlobalMemory path was selected, we should maybe rename it.
| TEST(RfTests, HighClassCountSplitHistogramFallsBackToGlobalMemory) | |
| TEST(RfTests, HighClassCountSplitHistogram) |
|
/merge |
Summary
Closes #8274.
This adds a runtime fallback for the RF/DT split histogram path when the per-block shared-memory histogram is too large or likely to reduce occupancy too much.
Instead of failing with the shared-memory limit for large
n_classes * max_n_bins,computeSplitKernelnow chooses between:The fallback is selected with a runtime boolean and shared/global pointer switching in the same kernel, so we avoid a separate implementation of the split logic. The threshold is intentionally named/commented as a tunable performance heuristic; it currently switches away from shared memory when the dynamic histogram allocation exceeds 16 KiB.
Details
use_global_memory_histogramto the split kernel launcher.Benchmarks
Local RF split benchmarks on NVIDIA RTX PRO 6000 Blackwell, driver 580.159.03.
Shared-path overhead from adding the runtime selection was small: about 0.7-1.2% in the measured shared-memory case.
Forced global-memory histogram vs normal shared-memory path for workloads that fit in shared memory:
2 classes x 32 bins8 classes x 32 bins16 classes x 64 bins32 classes x 128 bins44 classes x 128 binsRegression/default-style case:
1 output x 128 binsThis is why the fallback uses a 16 KiB tunable threshold rather than switching all shared-fit problems to global memory.
Testing
git diff --checklibcumlbuildRfTests.HighClassCountSplitHistogramFallsBackToGlobalMemory