[STF] C API and Python bindings for placement evaluation and structured partitions#2
Open
caugonnet wants to merge 58 commits into
Open
[STF] C API and Python bindings for placement evaluation and structured partitions#2caugonnet wants to merge 58 commits into
caugonnet wants to merge 58 commits into
Conversation
…VIDIA#9758) The blocked_partition alias is re-exported into the stf namespace, but the underlying blocked_partition_custom template it is defined from was not, which broke user code selecting the partitioned dimension explicitly (e.g. blocked_partition_custom<1>). Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
…9720) * Make stream and memory resource explicit in hyperloglog APIs * Apply cudax cuco conventions to hyperloglog headers * Address review feedback: stream-ordered ctor docs, granular no_init include, buffer-based tests * Add #else/#endif comments and hide defaulted special members from ABI * Use to_address instead of subscripting the iterator in host code
* Proper designated init for ReducePolicy * Extend with sub policy tests * Rename policy tests * Use a single test per device algorithm * Test operator<<
* Strip zero-width or unprintable unicode characters
Co-authored-by: Bernhard Manfred Gruber <bernhardmgruber@gmail.com> Co-authored-by: Jacob Faibussowitsch <jacob.fai@gmail.com>
Fixes NVBug: 6437537
…it conversions (NVIDIA#9772) * Fully classify the InputValue template to avoid implicit conversions
…e examples with new wrapper that does not depend on cupy or numba-cuda (NVIDIA#9653)
* Handle a type of size exactly 3 in cuda::buffer
* Extend with sub policy tests * Deprecate all agent policies * Unify deprecation messages
* bugprone-branch-clone
* Add builds for MSVC cccl_c_parallel * Disable SEH exceptions for catch2 * [c.parallel.v2] Run host-JIT Clang frontend on an 8 MB stack (fix Windows stack overflow) The host JIT invokes clang::CompilerInstance::ExecuteAction directly, which bypasses the clang driver's runWithSufficientStackSpace() guard. On Windows the default main-thread stack is 1 MB, so Clang's deep frontend recursion (recursive-descent parsing / template instantiation) overflows the stack while compiling heavier kernels such as radix_sort and segmented_reduce. On Linux the 8 MB default stack hides it. Run each ExecuteAction on a worker thread whose stack matches Clang's own DesiredStackSize (8 MB) -- the value already proven sufficient on Linux. * Drop override * Fix formatting --------- Co-authored-by: romanso <3954220+romanso@users.noreply.github.com>
Co-authored-by: Cédric Augonnet <158148890+caugonnet@users.noreply.github.com>
* Use cccl runtime in thrust tests pt 2 * Avoid make_pinned_buffer in Runtime test helpers * Assume device 0 for Runtime host test buffers * Hoist Runtime random buffer generators
cyclic_partition::get_executor was left as abort(), making cyclic the only partitioner unusable as a composite data place mapper, which needs the inverse coordinate-to-place direction. Implement the exact inverse of apply for zero-based coordinates: the owner of an element is its coordinate modulo the grid extent, independently in each dimension. Add a co-located unittest sweeping every place of a 2x3 grid over a 3-D box and checking that each coordinate visited by apply maps back to the place that produced it, covering dimensions the grid does not split. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Composite data places distribute an allocation according to a
partitioner that maps element coordinates to places, but the raw
allocation entry only carried a byte count: the partitioner received
flat byte offsets with fabricated 1-D extents, so any coordinate-based
mapper silently placed everything on one device. The tensor geometry
must reach the allocation seam.
Add a second allocate(data_dims, elemsize, stream) virtual to
data_place_interface with a base default forwarding to the byte-count
form: regular places inherit it unchanged ("same code, different
places" now extends to geometry-aware callers), while composite places
override it with the real localized_array path and reject the
byte-count form, whose implicit 1-D mode is removed. dim4 gains
index_to_pos(), the inverse of get_index(), shared by the runtime
slice interfaces and the raw composite path so the two linearization
conventions cannot drift.
Also make the block-owner decision reproducible (one seeded generator
per placement instead of a fresh random_device per block, probe count
now a parameter) and return placement statistics as a struct
(localized_stats: bytes per place, sampling accuracy - previously
stderr-only under CUDASTF_LOCALIZED_ALLOC_STATS, whose output is
preserved). A new evaluate_localized_placement() entry runs the exact
same decision procedure without allocating, so a candidate mapping can
be scored and tuned before committing memory.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Add a structured way to express how a tensor is distributed over a
grid of places: a per-dimension specification in the style of JAX's
PartitionSpec ("dimension 1, blocked over grid axis 0") lowered by
make_partition() onto a CuTe-style two-mode layout of flattened
(extent, stride) leaves.
Split dimensions are padded up to divisibility - the CuTe predication
idiom (partition the rounded shape, predicate against the true
extents) - which makes the layout exact and bijective over the padded
space. Exactness is what keeps the machinery trivial: validation is an
O(leaves) sorted-stride mixed-radix check, ownership queries are a
closed-form divmod chain, and no CUTLASS/CuTe dependency is needed
since only that subset of the layout algebra is exercised. Accessors
(dims, leaves, per-place offsets) expose everything a front-end needs
to derive further artifacts from the same object that drives data
placement.
The partition feeds the existing localized allocation machinery as a
stateful owner function: evaluate_localized_placement() accepts it
directly, and make_composite_data_place() wraps it as a composite data
place (per-tensor by nature, since a padded partition is specific to
one tensor's extents; the block-majority engine still decides where
each block lives). Routing such places through the runtime logical
data path is a recorded follow-up.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
placement.cu covers the evaluator (block-aligned split, majority tie-breaking on a straddling block, seeded determinism, cute partition agreeing with the equivalent classic partitioner), geometry-aware allocation on both composite flavors including kernel write/readback, and the rejection paths (byte-count allocate on a composite, extent mismatch with a partition). With two or more devices it additionally asserts the physical residency of each block on its owner and exercises the peer-access mappings; those checks self-skip on single-GPU machines so the multi-GPU CI provides the real coverage. cute_partition.cuh carries co-located unittests for the builders (including the uneven-extent padding case that motivates the padded design), owner sweeps, and validation rejections. The partitioned_axpy example shows the intended workflow: express the partition once, score it with evaluate_localized_placement() before allocating, run STF tasks over logical data on a composite place, and perform a raw geometry-aware allocation outside any STF context. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The bytes-per-place map is keyed by place strings, which cannot cross a C FFI boundary usefully. Also record bytes per linear grid index so bindings can return the distribution as a plain array. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
# Conflicts: # cudax/include/cuda/experimental/places.cuh
Expose the placement facilities of the places layer through the C API: - stf_data_place_allocate_shaped(): the geometry-aware allocation entry (any place handle; regular places use the byte-count path, composite places place each block on its owner - the plain byte-count stf_data_place_allocate() now fails on composites since a byte count alone cannot carry the tensor geometry the partitioner needs). - stf_placement_evaluate() / stf_placement_evaluate_partition(): score a candidate mapping without allocating, returning stf_placement_stats plus optional per-grid-position byte counts. - stf_cute_partition_*: build structured partitions from a JAX-like per-dimension spec (or raw leaves), with accessors exposing dims, leaves and per-place offsets so bindings can derive further artifacts from the same object; stf_data_place_composite_cute() wraps one as a composite data place. - stf_partition_fn_blocked()/stf_partition_fn_cyclic(): native partition functions usable wherever a mapper is expected, avoiding FFI callback cost for the common policies. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Python surface over the new placement C API:
- cute_partition with from_spec() (JAX-like per-dimension entries such
as ("blocked", axis)) and from_leaves(), plus accessor properties
(dims, leaves, place offsets) for front-ends deriving launch-side
artifacts from the same object that drives data placement.
- placement_evaluate() scoring a mapping before any allocation; the
mapper can be a cute_partition, a native partition function pointer
(partition_fn_blocked/partition_fn_cyclic, no FFI callback cost), or
a Python callable (which crosses the GIL per probe - documented as
the slow generic path). Returns placement_stats with accuracy and
per-grid-position bytes.
- data_place.allocate() now accepts extents (with keyword-only
elemsize) in addition to a byte count; composite places require the
extents form. data_place.composite() additionally accepts native
partition function pointers; data_place.composite_cute() wraps a
structured partition.
The pytest suite covers builders, all three mapper forms agreeing, the
seeded majority tie-breaking case, shaped allocation on both composite
flavors, and a multi-GPU residency check that self-skips below two
devices (real coverage comes from multi-GPU CI).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Overloading allocate() made every data_place_interface subclass that overrides the byte-count form partially override the virtual set, tripping -Woverloaded-virtual until it adds a using-declaration - a recurring tax on every current and future place implementation. A distinct name removes the whole problem class (no name hiding, no overload-ambiguity concerns) and aligns with the C API, which cannot overload and needs a distinct name anyway. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The geometry-aware allocation entry is stf_data_place_allocate_nd(), matching the C++ data_place::allocate_nd() it wraps. The Python data_place.allocate() keeps dispatching on the argument type (byte count vs extents), which is the Python idiom, and now calls the renamed entry underneath. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
cyclic_partition::get_executor was left as abort(), making cyclic the only partitioner unusable as a composite data place mapper, which needs the inverse coordinate-to-place direction. Implement the exact inverse of apply for zero-based coordinates: the owner of an element is its coordinate modulo the grid extent, independently in each dimension. Add a co-located unittest sweeping every place of a 2x3 grid over a 3-D box and checking that each coordinate visited by apply maps back to the place that produced it, covering dimensions the grid does not split. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Adversarial review fixes, none changing the intended semantics: - Guard rank-0 (scalar) slices in the composite delinearize: their data_dims report an extent of 0, so index_to_pos would divide by zero where the previous per-rank lambda returned the origin. - Validate elemsize (>= 1, <= block size) at the placement seams instead of dividing by zero or silently sampling garbage when a block holds no whole element. - Compute each block's first element exactly (i * block_size / elemsize, as before the refactor) so non-dividing element sizes do not accumulate a sampling-window drift across blocks. - Check the partition's grid extents against the execution place grid (cute composite construction and evaluation): a mismatched pair previously mis-attributed bytes or aborted deep in get_place. - Reject zero-extent leaves and grid axes, and out-of-range place_offset indices, instead of wrapping silently. - Compare data_dims for real in the allocation-cache equality (was assert-only: transposed extents could produce a false cache hit in release builds). - Make the straddling-block test structural (62.5/37.5 split: the all-probes-in-majority probability is ~1e-13 for any uniform sampler) instead of relying on the stdlib's implementation-defined distribution sequence. - Fall back to a single device in the partitioned AXPY example when peer access is unavailable (the single task touches every page from each device), and factor the default block-size selection into one helper. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Adversarial FFI review fixes; every crash below was reproduced from pure Python before the fix: - cute_partition() direct instantiation left a NULL handle that accessors dereferenced (SIGSEGV): __init__ now raises, pointing at from_spec()/from_leaves(). - bool passed as a mapper matched isinstance(int) and was called as a function pointer (SIGSEGV): bools and NULL pointers are rejected in placement_evaluate() and data_place.composite(). - A Python mapper that raised was swallowed by the ctypes trampoline and evaluation reported success with wrong statistics: trampoline errors are now recorded and re-raised (chained) by placement_evaluate, and printed for allocation-time callbacks. - stf_fill_placement_outputs wrote the caller's bytes_per_grid_index buffer at unvalidated mapper-produced keys (heap corruption with a scalar exec place): keys are validated against the grid size and the evaluation fails cleanly; the helper also moves out of the extern "C" block. - stf_cute_partition_place_offset now catches the out-of-range throw (returns UINT64_MAX, documented) instead of terminating; from_leaves gained null-array asserts; the byte-count allocate() rejects a non-default elemsize instead of ignoring it; sequence inputs that yield more items than their length report no longer overflow the fixed marshalling arrays; evaluating a cute partition against mismatched data_dims is rejected. - The multi-GPU residency test queries the device's real allocation granularity instead of assuming 2 MiB, the straddling-block test uses a structural 62.5/37.5 geometry, and a new pytest pins every crash path above as a clean Python exception. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
A hand-built dim4 can carry zero extents, which would make the modulo undefined; real grids already assert positive dimensions, so this only guards direct misuse. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
# Conflicts: # cudax/include/cuda/experimental/__places/partitions/cyclic_shape.cuh
localized_array already validates VMM support with a clear error, and no other composite-place example performs its own check; the raw driver boilerplate distracted from the example's point. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ample exec_place::all_devices() already enumerates the devices, and the places machine sets up peer and mempool access between them - the example re-implemented both, plus a peer-access fallback that only existed because it launched a single kernel over the whole composite buffer from one device. Follow the idiomatic grid-task pattern of 01-axpy-places instead (each place computes its own blocked chunk on its own stream), which needs no cross-device access at all. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
With ceil-division chunks, i * chunk can exceed N when there are more places than chunks; the unsigned N - start then underflows and the kernel writes out of bounds. Unreachable with the example's sizes, but the loop demonstrates the canonical pattern and must be safe in general. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The example's memory resource byte-allocated on a composite place, which now throws (a byte count alone cannot carry the geometry the partitioner needs). A memory resource genuinely hands out untyped bytes, so it states that explicitly with allocate_nd(dim4(bytes), 1) - equivalent for every other place type, and byte-granular distribution for composites, matching the previous behavior. Document allocate_nd and the composite requirement in the places guide. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Replace the std::vector leaf storage of cute_partition with cuda::std::array (bounded by a documented max_leaves = 16 per mode; make_partition emits at most 2 per dimension). The partition becomes trivially copyable and its queries (owner, dims, leaf spans) are host/device callable, so it can cross the kernel boundary by value - the property a future parallel_for integration needs, with no separate device-side snapshot type. Accessors now return cuda::std::span over the filled prefix; construction keeps its std::vector convenience signature and rejects leaf counts beyond the capacity. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
14 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Python/C surface for the placement facilities of NVIDIA#9804, stacked on that PR and NVIDIA#9803 (both merged into this branch; their commits appear in the diff until they land on main and flow into
stf_c_api).C API (
cccl/c/experimental/stf/stf.h):stf_data_place_allocate_nd()— geometry-aware allocation on any place handle.stf_placement_evaluate()/stf_placement_evaluate_partition()— score a candidate mapping without allocating:stf_placement_statsplus optional per-grid-position byte counts.stf_cute_partition_*— build structured partitions from a JAX-like per-dimension spec or raw leaves, with accessors (dims, leaves, per-place offsets);stf_data_place_composite_cute()wraps one as a composite data place.stf_partition_fn_blocked()/stf_partition_fn_cyclic()— native partition functions usable wherever a mapper is expected, avoiding FFI callback cost for the common policies.Python (
cuda.stf._experimental):cute_partition.from_spec((nx, ny, nz), (None, ("blocked", 0), None), (nplaces,))andfrom_leaves, with accessor properties.placement_evaluate(grid, mapper, dims, elemsize)where the mapper is acute_partition, a native function pointer, or a Python callable (documented as the slow path — each probe crosses the GIL); returnsplacement_statswithaccuracyand per-position bytes.data_place.allocate()accepts extents with keyword-onlyelemsizein addition to a byte count (existing callers unaffected —streamstays the second positional);composite()additionally accepts native function pointers;composite_cute()wraps a structured partition.Hardened by adversarial review with every crash reproduced first: NULL-handle accessors,
bool-as-function-pointer,elemsize=0SIGFPE, an out-of-range mapper corrupting the caller's stats buffer, and raising Python mappers being silently swallowed — all now clean Python exceptions, each pinned by a pytest. The multi-GPU residency test queries the device's real allocation granularity and self-skips below two devices.Checklist
🤖 Generated with Claude Code