Skip to content

feat(cutile-rs): add per-operation atomic aggregation - #1558

Open
robertream wants to merge 1 commit into
NVIDIA:mainfrom
robertream:feat/aggregation-per-operation
Open

robertream wants to merge 1 commit into
NVIDIA:mainfrom
robertream:feat/aggregation-per-operation

Conversation

@robertream

@robertream robertream commented Oct 20, 2026 •

Copy link
Copy Markdown

Adds typed per-operation atomic aggregation capabilities to cuTile: AddTensor, AndTensor, OrTensor, XorTensor, MaxTensor, and MinTensor.

Each capability exposes only its matching reduction, validates supported element types, enforces tiling and index bounds, rejects invalid aliases, and includes compile-only and GPU test coverage.

The implementation is compile-verified against cutile_ir::Module. All nine GPU tests passed against the aggregation implementation included in this PR.

Depends on PR #1557, which preserves full-width integer literal encodings required by unsigned max/min tests.


I used an AI assistant (Claude) to write and check this change. I reviewed the diff and validated it.

@copy-pr-bot

copy-pr-bot Bot commented Oct 20, 2026

Copy link
Copy Markdown

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.

@robertream
robertream force-pushed the feat/aggregation-per-operation branch 3 times, most recently from ee609da to 4b79c27 Compare October 20, 2026 05:22
@robertream
robertream marked this pull request as ready for review October 20, 2026 05:34
@robertream
robertream requested a review from elibol as a code owner October 20, 2026 05:34
@robertream
robertream force-pushed the feat/aggregation-per-operation branch 5 times, most recently from 02f84e7 to 087e532 Compare October 21, 2026 02:49
Adds typed atomic aggregation capabilities with compile-time operation and element validation, launch-time alias checks, and compile-only and GPU test coverage.

Depends on NVIDIA#1557 (fix/cutile-compiler: preserve unsigned integer literal ranges).

Storage identity is opaque: `Tensor::storage_id`, the stored-binding
traits, and `validate_binding_aliases` carry a `StorageId` (equality is its
whole public surface) rather than a bare `usize` address.

Library kernels built on these capabilities live in `cutile-kernels`,
generic over the element wherever the added value comes from an input
(`scatter_add`, `embedding_backward_scatter_add`, `segment_sum`,
`weighted_histogram` over `AtomicAdd`; `max_reduce` over `AtomicOrd` with the
element selecting signed vs unsigned for both the intra-tile reduction and
the atomic; `bitmask_or` over `AtomicBitwise`).
Compile-only tests drive each kernel across its whole element set asserting
the emitted mode per pair, host references cover the edge cases, and device
tests compare against them.

GPU correctness of the unsigned-comparison and bf16/64-bit-constant cells
depends on the literal-encoding fixes (PR NVIDIA#1557); the combined tree passed
20/20 required tests on an H100 (sm_90).

`ElementType` gains a `ONE` constant, so every kernel is generic over
its element, `histogram` included (its increment is `E::ONE` materialized
on-device).

Signed-off-by: Robert Ream <154010+robertream@users.noreply.github.com>
@robertream
robertream force-pushed the feat/aggregation-per-operation branch from 087e532 to 6fdc67e Compare October 21, 2026 02:59

This branch has not been deployed

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant