Optimize partition - #25154

Description

Optimize cudf::partition by selecting partition-map construction and result materialization from cached CPU metadata. The implementation reduces atomic contention, intermediate-map traffic, temporary memory, kernel-launch overhead and host/stream synchronization, while materializing eligible columns together.

The heuristics use row count, partition count, column count, fixed-width storage size, cached null counts, STRING offset type, and the character-buffer pointer. Selection is O(number of columns), with no GPU sampling, length/skew measurement, or device readback to select a performance plan. Mandatory STRING allocation-bound readbacks remain separate from heuristic selection. A validity mask with null_count() == 0 qualifies for non-null paths.

All implemented optimizations

1. Remove per-row partition-offset atomicAdd from optimized rank/scan paths. Warp-local histograms use matching-lane ballots, population counts, and one elected writer per bin. A partition-major prefix scan supplies partition bases and preceding-block counts together, replacing contended per-row global cursor increments.

2. Specialize two-partition work. CUB DevicePartition::If uses a counting input iterator and an output proxy to emit directly to the selected sink, avoiding a separate materialized row-index input and unnecessary general-bin machinery.

3. Save compact local ranks when recomputation is more expensive. Tile-local ranks fit in uint16_t (2 bytes per row), rather than a full-width rank. The selected saved-rank path reuses them during emission.

4. Recompute ranks to eliminate the N-entry saved-rank buffer. Histogram-only counting followed by rank recomputation trades extra map reads for less temporary memory and global rank-buffer traffic.

5. Use two-tile superblocks where selected. The super2 gather-map path handles two adjacent 1,024-row microtiles per block, carries partition bases between them, and reduces block-count/prefix storage without storing per-row ranks globally.

6. Amortize global reservation atomics. The reserve plan computes local ranks and makes one global cursor reservation per nonempty block/bin, rather than one per row. This plan still uses atomicAdd; nullable bitmap updates and the generic fallback also retain atomics.

7. Write eligible fixed-width result columns directly. Fuse destination calculation with payload and validity emission, removing the N-entry gather map and the subsequent generic gather pass. A shared table sink materializes all selected fixed-width columns in the same kernel invocation instead of separate per-column gather launches; 1/2/4/8/16-byte storage has specialized aligned copies.

8. Reorder wide output within shared-memory tiles. Group source rows/destinations by partition before emission to improve destination-store locality. This preserves direct output and avoids a global row permutation.

9. Exploit cached null counts. Skip invalid payload reads/writes, omit output masks for zero-null columns, and reuse the input null counts because a complete permutation preserves them. Direct materialization bypasses generic validity-count calculation/readback, saving that host synchronization and associated stream.sync waiting time.

10. Bypass permutation for identity cases. One partition or one row uses direct copies or an implicit identity iterator rather than an N-entry row map. Fixed-width data uses device copies; non-null STRING identity copies reuse/rebase offsets and copy the selected character range directly. Nullable STRING/nested values retain generic identity gather.

11. Fold uniform all-NULL and CPU-known empty-STRING columns. Exclude invariant columns from partition work, partition only the reduced schema, then restore them after internal row-map/rank temporaries are released. Entirely uniform tables need boundaries only, not a row permutation; zero-column tables with rows follow the same boundary-only handling.

12. Calculate STRING allocation sizes across result columns together. One sizing kernel reads every eligible column's sliced boundary offsets, followed by one batched pinned-host size readback. A permutation preserves each column's total character count, so output allocation does not require a separate post-gather size readback for every column. Identity STRING bounds are also read in one batch.

13. Save synchronization time by batching STRING preparation. Allocate character buffers from the batched input-size readback, enqueue output-offset scans using reusable CUB scratch storage, and construct result columns without sequential per-column allocation-size readbacks. GPU operations remain ordered on the supplied stream; scans are still per column, not concurrent-stream scans. Required allocation/type-limit checks, descriptor uploads, and host partition-boundary readbacks can still synchronize. No isolated millisecond saving for synchronization is claimed; existing measurements are end-to-end.

14. Fuse gathering STRING columns into one character-copy kernel launch. A two-dimensional grid assigns columns through blockIdx.y, allowing blocks for multiple result columns to execute concurrently, subject to GPU scheduling. The general fused path shares the gather/inverse maps; other selected STRING paths scatter all column sizes together. This avoids one character-copy launch per result column. Offset scans and allocation preparation remain separate stages.

15. Reduce row maps for two-partition, small-schema STRING inputs. Below 4M rows, retain one scatter map and emit output lengths during partitioning. At/above 4M rows, a single STRING column recomputes destinations into a 1,024-entry shared tile with no N-entry global row map; two/three columns store the binary bin plus tile-local rank in one uint16_t per row, avoiding full-width maps and repeated partition-map reads during character copy.

16. Reduce STRING copy register lifetime and dependent loads. The mapless path stores destinations in shared memory and keeps its copy loop rolled so rank/key registers can die before variable-length copying. The packed-rank path loads both bin bases independently of the packed rank to allow prefix/rank load overlap.

17. Balance per-string copy work inside the selected kernel. Four-lane groups handle short strings; eligible rebalanced kernels use two 16-lane groups for medium strings; full warps handle long strings with vector loads and aligned 16-byte stores. Boundary bytes use scalar copies, with loads bounded to the individual string. This avoids a separate length-sampling or host dispatch pass.

18. Avoid empty and duplicate preparation work. Zero-character columns use zeroed offsets and skip character copy; fused STRING scans reuse one scratch allocation. Generic gather-map construction uses histogram-only/two-tile kernels, and supported rank/gather-map helpers read host boundaries through pinned staging.

19. Remove redundant block barriers and shared histogram code. Histogram-only tiles retain the barrier before shared-memory reuse; reordered emission relies on its caller's synchronization. Saved local ranks reuse the existing warp histogram helper. Identity-copy headers declare their dependencies, and trivial partition boundaries use pinned D2H staging.

All heuristic scenarios

The following tables describe the current dispatch order and thresholds. Let N be rows, P partitions, C columns, W total fixed-width bytes per row, F mean all columns are fixed-width, and A mean at least one column has an actual null (null_count() > 0). After removing invariant columns, C/W/A refer to the reduced schema. “Large” means N >= 1,000,000. “Dense-null” means F, A, large, and every column is at least 75% null. direct emits result payloads immediately; otherwise the sink builds a gather map for generic gather. Rows within each table are considered in order unless marked as an override.

1. Empty, identity, and invariant inputs (before algorithm selection)

2. Specialized STRING plans (before fixed-width/general selection)

Admission requires N >= 1,024, 2 <= P <= 256, 1 <= C <= 65,535, all columns STRING, zero actual nulls, and INT32 input offsets. P > 2 additionally requires C >= 2. Each output character-buffer size must be below the configured INT64-offset threshold; otherwise this specialization declines and dispatch continues to generic handling.

3. Fixed-width/general partition-plan selection

Admission requires N >= 1,024 and 2 <= P <= 256. STRING admission above has already been attempted.

4. Block-rank and generic fallbacks

5. Device-side STRING copy work sharing (not CPU plan selection)

These are measured heuristic tradeoffs for CPU-metadata groups, not a guarantee of the fastest algorithm for every map skew or string-length distribution. The API, sliced offsets, offset-width limits, supplied stream, output resource, and generic fallback semantics remain unchanged.

Validation

Current tested head on NVIDIA L4 (sm89), CUDA 12.9.1 / GCC 14.2.1, Release, per-thread default stream:

• PARTITIONING_TEST: 2,140 passed, one existing disabled test; STREAM_PARTITIONING_TEST: 7 passed.
• CTest stream-interposition run passed (including its resource fixture).
• Compute Sanitizer: 15 memcheck cases and 5 synccheck cases, with zero errors; 5 targeted racecheck cases, with zero errors or warnings. Racecheck covers the modified histogram, rank and reordered-emission kernels.
• Coverage includes 48 exact host-oracle cases with 1,000,003 rows each (partial final tile, data/map slices, 1/2/4/8 INT64 columns, P2/P100, null density and skew), single/two-column large sparse STRING copies, map types, offset limits, mixed/nested schemas, non-default streams and resource handling.
• The identity header compiles independently with the production CUDA flags. Formatting, spelling and whitespace checks passed. Redundant block barriers were checked in generated sm89 SASS: one/two-tile histogram 3->2 / 6->4 sites; reordered output 10->9 sites.

Performance and memory

Fresh baseline-versus-PR measurement, 2026-11-06.

• Baseline: PR base commit 65b9b3c222ee7d6ba8f0f442688bd64268709e93.
• PR: tested head c285951eeddfae0a34379d82f0c40dbbd264e695.
• Hardware/build: NVIDIA L4 (sm89, 23,034 MiB reported VRAM), CUDA 12.9.1, GCC 14.2.1, Release (-O3, C++20), per-thread default stream, RMM CUDA async allocator.
• The same frozen native cudf::partition benchmark and 348 identical fixtures were run freshly against both exact libraries. Each case has a minimum 0.15-second warmup and 15 timed samples. Input construction and independent per-partition row-hash multiset checks are outside timing.
• CUDA-event elapsed time spans the API call, including stream gaps caused by required host work. Wall time includes allocations/readbacks, output destruction and final event synchronization. These are native cuDF measurements, not Spark/JNI end-to-end timings.
• Every initial >5% slowdown candidate in either timing metric, plus key cases, was rerun: 15 fixtures, three paired alternating-order rounds, 30 samples per case per library. All accepted runs passed correctness, library/executable/fixture hash and competing-compute-process checks.

Speedups are equal-fixture-weight geometric means of per-case median baseline/PR time. They are descriptive results, not workload-weighted predictions or confidence intervals. GPU peak is tracked output + temporary allocations, excluding input, context/driver/code caches and host pinned staging. Peak reduction is the arithmetic mean of each case's reduction, not the reduction of a pooled memory total.

Breakdown versus the same baseline. Large means N >= 1M with 2 <= P <= 256; each group includes its selected/fallback behavior. NULL percentages are fixture-generation targets; dispatch uses each column's actual cached null count.

Selected absolute timings and memory. Times are medians of the three 30-sample round medians; the displayed speedup is the ratio of those displayed timing medians. MiB values are output + temporary GPU allocation peaks.

Regression check versus baseline. A confirmed regression requires >5% slowdown in every paired repeat round for the same metric. This does not claim every individual case is faster.

None of the initial >5% slowdown candidates remained >5% slower in all three rounds (15 fixtures repeated). Smaller or inconsistent slowdowns may still exist; see the representative fallback case above.

CPU-metadata heuristics make speed/memory tradeoffs across inputs sharing the same schema metadata; they do not use observed string lengths or map skew to dispatch. No physical DRAM transaction measurement is claimed.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.