Skip to content

Select: load only the flagged items when the selection is dense - #11503

Open
0z5a wants to merge 1 commit into
NVIDIA:mainfrom
0z5a:select-if-predicated-item-loads
Open

0z5a wants to merge 1 commit into
NVIDIA:mainfrom
0z5a:select-if-predicated-item-loads

Conversation

@0z5a

@0z5a 0z5a commented Sep 18, 2026

Copy link
Copy Markdown

DeviceSelect with USE_SELECT_FLAGS re-reads the whole input tile after the flag scan, so an
invocation that keeps a small fraction of the items still pays full HBM traffic for the payload
(issue #326). When the predicate is a per-item flag and the payload is trivially copyable, the tile
loop can read only the surviving indices and copy them straight out.

The fast path is gated on SELECT_METHOD == USE_SELECT_FLAGS && SelectionOpt == SelectImpl::Select && is_trivially_copyable_v<InputT>, so every other configuration keeps the existing code path.

Measured on an L20 (sm89, CUDA 12.8), ABBA-paired against the unpatched agent: 2.26x-3.48x at a
selection rate of 0.1% or below, parity from 50% upwards where the flag scan dominates. Correctness:
1040/1040 flagged-select cases and 120/120 shared-agent regression cases.

What the patch adds

  • can_predicate_item_loads — a compile-time gate that turns the fast path on only for
    SELECT_METHOD == USE_SELECT_FLAGS, SelectionOpt == SelectImpl::Select and a trivially copyable
    payload, so every other configuration keeps the existing code path untouched;
  • LoadItemsPredicated — reads only the flagged indices out of the tile instead of re-reading the
    whole tile, and copies survivors straight to the output;
  • a new first branch in PrepareTile that selects between the two.

Correctness on the same machine: 1040/1040 flagged-select cases and 120/120 shared-agent regression
cases.

Measurement environment

The numbers above were taken on an L20 (sm89) with driver 570.86.10 / CUDA 12.8. The machine has
since been upgraded to driver 595.91.07, so the figures are being re-measured and this section
will carry the updated table; treat the current numbers as driver-570 references.

DeviceSelect with USE_SELECT_FLAGS re-reads the whole input tile after the flag
scan, so an invocation that keeps a small fraction of the items still pays full HBM
traffic for the payload. When the predicate is a per-item flag and the payload is
trivially copyable, the tile loop can instead read only the surviving indices and
copy them straight to the output.

Measured on an L20 (sm89, CUDA 12.8), ABBA-paired against the unpatched agent:
2.26x-3.48x when the selection rate is 0.1% or below, and parity from 50% upwards,
where the flag scan already dominates. Correctness: 1040/1040 flagged-select cases
and 120/120 shared-agent regression cases pass.
@copy-pr-bot

copy-pr-bot Bot commented Sep 18, 2026

Copy link
Copy Markdown
Contributor

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@github-project-automation github-project-automation Bot moved this to Todo in CCCL Sep 18, 2026
@cccl-authenticator-app cccl-authenticator-app Bot moved this from Todo to In Progress in CCCL Sep 18, 2026
@0z5a
0z5a marked this pull request as ready for review September 19, 2026 01:40
@0z5a
0z5a requested a review from a team as a code owner September 19, 2026 01:40
@0z5a
0z5a requested a review from davebayer September 19, 2026 01:40
@cccl-authenticator-app cccl-authenticator-app Bot moved this from In Progress to In Review in CCCL Sep 19, 2026
@coderabbitai

coderabbitai Bot commented Sep 19, 2026

Copy link
Copy Markdown
Contributor

Review Change StackReview Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Repository: NVIDIA/cccl/.coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 5d80bbc4-48b4-4873-be6f-552b91b3dc72

📥 Commits

Reviewing files that changed from the base of the PR and between 2fcc3f0 and bdb0dbb.

📒 Files selected for processing (1)
  • cub/cub/agent/agent_select_if.cuh

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.


📝 Summary

Summary by CodeRabbit

  • Performance
    • Improved flagged selection performance by loading payload data only for items that are selected.
    • Reduced unnecessary data movement for eligible trivially copyable inputs.
    • Preserved correct selection counts and handling of incomplete final tiles.

Walkthrough

Changes

AgentSelectIf now supports predicated payload loading for eligible pure flagged selection. It initializes flags before payload loading, loads only selected items, and clamps out-of-range final-tile accesses.

Select-If loading path

Layer / File(s) Summary
Eligibility and predicated tile loading
cub/cub/agent/agent_select_if.cuh
Adds can_predicate_item_loads and LoadItemsPredicated<IS_LAST_TILE>. Eligible selections load flagged items only. Final-tile loads use a valid element for out-of-range slots. Other selection modes keep their existing loading paths.

Suggested reviewers: bernhardmgruber

Priority: ➖ Normal

Change: Refactor

✨ Finishing Touches 💡 1
🛠️ Fix failing CI checks 💡
  • Commit to this branch
  • Create a new PR

Comment @coderabbitai help to get the list of available commands.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Status: In Review

Development

Successfully merging this pull request may close these issues.

1 participant