[STF] Expose geometry-aware allocation and native partition functions in the C API#10038
Conversation
… in the C API Make composite data places usable end to end from C: - stf_partition_fn_blocked(dim) / stf_partition_fn_cyclic() expose the existing blocked/cyclic partitioners as native stf_get_executor_fn values, so composite places can be built without FFI callback cost. - stf_data_place_allocate_nd(dims, elemsize) forwards to the C++ data_place::allocate_nd already on main: equivalent to a byte-count allocation on ordinary places, geometry-aware block placement on composite ones (where the plain byte-count stf_data_place_allocate correctly fails, since a byte count alone cannot carry the geometry). Both additions only use types already in the C API (stf_dim4, stf_get_executor_fn); nothing here depends on structured (cute) partitions, which remain in a separate PR. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughChangesThe STF C API adds blocked and cyclic partition helpers and tensor-aware ND allocation. Implementations validate shaped allocation inputs and blocked partition geometry, while tests cover ordinary and composite placements, partition dimensions, device round trips, and overflow rejection. Suggested reviewers: Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 76d91229-5f74-42e1-9656-f1dc8323c749
📒 Files selected for processing (3)
c/experimental/stf/include/cccl/c/experimental/stf/stf.hc/experimental/stf/src/stf.cuc/experimental/stf/test/test_allocate_nd.cpp
- Rename MiB to one_mib (snake_case per coding guidelines). - Declare immutable locals const throughout. - Cover every blocked dimension selector (0-3) and out-of-range values (-1, 4) with a 4D shape divisible by the grid along each axis. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 1
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 9418e47f-0cb0-4642-992e-34f81a1a9095
📒 Files selected for processing (1)
c/experimental/stf/test/test_allocate_nd.cpp
Compile-time test constants (one_mib, shapes, element counts) are now constexpr, with the namespace-scope constant inline constexpr. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
/ok to test 66a3a03 |
This comment has been minimized.
This comment has been minimized.
| ::std::memcpy(&dims, data_dims, sizeof(dims)); | ||
| try | ||
| { | ||
| return from_opaque(h)->allocate_nd(dims, elemsize, stream); |
There was a problem hiding this comment.
Important: Validate the extent and byte-count multiplications before forwarding. With dims={UINT64_MAX, UINT64_MAX, 1, 1} and elemsize=1, dim4::size() wraps to 1: an ordinary place returns a non-null one-byte allocation, while a two-place blocked composite terminates with SIGFPE after the mapper computes a zero part_size. The try/catch cannot catch this. Please checked-multiply all extents and elemsize, return nullptr on overflow, and add regression coverage for both paths.
There was a problem hiding this comment.
Fixed in 8f38c13. data_place::allocate_nd() now checked-multiplies all four extents and elemsize, rejecting geometries whose byte count overflows size_t or exceeds PTRDIFF_MAX (std::invalid_argument, surfaced as NULL through the C API). The check sits at the single entry point, before any partitioner math runs, so both paths you describe are covered: the ordinary place no longer silently returns a wrapped (tiny) allocation, and the composite blocked place can no longer reach the part_size == 0 division (the SIGFPE). blocked_partition::get_executor additionally asserts nplaces > 0 and part_size > 0 to document the precondition. Regression coverage added on both paths, in C (test_allocate_nd.cpp, "shaped allocation rejects overflowing geometries": extent wrap, elemsize wrap, above-PTRDIFF_MAX, and the composite case that used to SIGFPE) and C++ (placement.cu, test_shaped_alloc_overflow).
| //! \brief Native blocked partition function for a given dimension | ||
| //! (values outside [0, 3] select the highest-rank dimension, like -1), | ||
| //! usable wherever an stf_get_executor_fn is expected without any FFI | ||
| //! callback cost. |
There was a problem hiding this comment.
Suggestion: document the clamping behavior to the highest axis whose extent is greater than one
There was a problem hiding this comment.
Done in 8f38c13. The doc now states that the requested dimension is clamped to the highest axis whose extent is greater than one — including the case where an in-range dim beyond that axis is clamped down (e.g. dim 2 on extents {n, 1, 1, 1} partitions along axis 0). The same clarification was added to the blocked_partition doc block in the C++ header.
…locked-partition clamping Address review feedback on NVIDIA#10038: - data_place::allocate_nd() now checked-multiplies the extents and element size, rejecting geometries whose byte count overflows size_t or exceeds PTRDIFF_MAX (std::invalid_argument, surfaced as NULL through the C API). Previously the product silently wrapped: an ordinary place returned a tiny allocation for an astronomically large tensor, and a composite blocked place died with SIGFPE on a zero part_size. - Assert the (now guaranteed) preconditions in blocked_partition::get_executor instead of dividing by zero. - Document that the blocked partition dimension is clamped to the highest axis whose extent is greater than one, including for in-range dimension selectors. - Add regression tests on both the C and C++ paths. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
cudax/include/cuda/experimental/__places/partitions/blocked_partition.cuh (1)
117-118: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winsuggestion:
nplacesis never reassigned after line 117; declare itconst. As per coding guidelines, "All variables that are not modified must be declared const, including cast results, function return values, and loop-invariant computations."♻️ Proposed fix
- size_t nplaces = grid_dims.x; + const size_t nplaces = grid_dims.x;Source: Coding guidelines
cudax/test/places/placement.cu (1)
272-282: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winsuggestion:
dev,grid, andcare never reassigned after initialization; declare themconst. As per coding guidelines, "All variables that are not modified must be declared const, including cast results, function return values, and loop-invariant computations."♻️ Proposed fix
- data_place dev = data_place::device(0); + const data_place dev = data_place::device(0); expect_invalid(dev, wrapping_dims, 1); ... - auto grid = make_device_grid(ndevs, 2); - data_place c = data_place::composite(blocked_partition_custom<1>{}, grid); + const auto grid = make_device_grid(ndevs, 2); + const data_place c = data_place::composite(blocked_partition_custom<1>{}, grid);Source: Coding guidelines
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d934ecbb-8252-4e46-ab0b-b12fa2d9b173
📒 Files selected for processing (5)
c/experimental/stf/include/cccl/c/experimental/stf/stf.hc/experimental/stf/test/test_allocate_nd.cppcudax/include/cuda/experimental/__places/partitions/blocked_partition.cuhcudax/include/cuda/experimental/__places/places.cuhcudax/test/places/placement.cu
🚧 Files skipped from review as they are similar to previous changes (1)
- c/experimental/stf/include/cccl/c/experimental/stf/stf.h
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
/ok to test 97b06e4 |
🥳 CI Workflow Results🟩 Finished in 1h 30m: Pass: 100%/61 | Total: 1d 18h | Max: 1h 29m | Hits: 6%/317954See results here. |
3 participants
Description
Make composite data places usable end to end from the C API:
stf_partition_fn_blocked(dim)/stf_partition_fn_cyclic()expose the existing blocked/cyclic partitioners as nativestf_get_executor_fnvalues, so C callers can build composite data places (stf_data_place_composite(), already on main) without hand-writing an FFI callback for the standard mappings.stf_data_place_allocate_nd(dims, elemsize)forwards to the C++data_place::allocate_ndalready on main: equivalent to a byte-count allocation on ordinary places, geometry-aware block placement on composite ones — where the plain byte-countstf_data_place_allocatecorrectly fails, since a byte count alone cannot carry the tensor geometry.Both additions only use types already present in the C API (
stf_dim4,stf_get_executor_fn). Nothing here depends on structured (cute) partitions; those, together with placement evaluation and the Python bindings, remain in #9892, which this PR is split out of so it can targetmaindirectly instead of being stacked on #5315.Tested with a new
test_allocate_nd.cpp: shaped allocation on an ordinary device place, clean byte-count failure plus successfulallocate_ndon composite places built with both native partition functions, with device round-trip verification.Checklist
🤖 Generated with Claude Code