[STF] Drive parallel_for and data placement from structured partitions#9808
Merged
Conversation
Contributor
2 tasks
Contributor
|
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 PR refactors CuTe partitions into typed partitions and runtime descriptors, integrates stateful partitioning and predicate-aware regions into STF Structured partition execution
Possibly related PRs
Suggested labels: Suggested reviewers: Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 8af68fc3-d748-4a5e-8ae9-14231140e485
📒 Files selected for processing (20)
cudax/examples/places/thrust_device_data_place_allocator.cucudax/examples/stf/CMakeLists.txtcudax/examples/stf/fdtd_mgpu.cucudax/examples/stf/partitioned_axpy.cucudax/include/cuda/experimental/__places/cute_partition.cuhcudax/include/cuda/experimental/__places/data_place_interface.cuhcudax/include/cuda/experimental/__places/localized_array.cuhcudax/include/cuda/experimental/__places/places.cuhcudax/include/cuda/experimental/__stf/graph/interfaces/slice.cuhcudax/include/cuda/experimental/__stf/internal/backend_ctx.cuhcudax/include/cuda/experimental/__stf/internal/parallel_for_scope.cuhcudax/include/cuda/experimental/__stf/internal/stf_places_extended_exports.cuhcudax/include/cuda/experimental/__stf/stream/interfaces/slice.cuhcudax/include/cuda/experimental/__stf/utility/dimensions.cuhcudax/include/cuda/experimental/places.cuhcudax/test/places/CMakeLists.txtcudax/test/places/placement.cucudax/test/stf/CMakeLists.txtcudax/test/stf/places/cute_parallel_for.cudocs/cudax/places.rst
Contributor
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: 1eb5f8a2-b5cd-4232-8645-4f38167b6452
📒 Files selected for processing (1)
docs/cudax/places.rst
Jacobfaib
approved these changes
Jul 14, 2026
caugonnet
force-pushed
the
stf-cute-parallel-for
branch
from
July 16, 2026 07:29
47c75f0 to
6a6fb2e
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
cudax/include/cuda/experimental/__stf/graph/interfaces/slice.cuh (2)
28-28: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winsuggestion: Add
<functional>directly. Line 124 uses::std::function; this header must not rely on transitive includes. As per coding guidelines, “Include all headers needed by the symbols being used; do not rely on transitive includes.”Source: Coding guidelines
101-131: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winsuggestion: Declare unmodified locals
const, includingdelinearize,cute_place,part, andcute_base. As per coding guidelines, “All variables that are not modified must be declaredconst.”Source: Coding guidelines
cudax/test/stf/CMakeLists.txt (1)
82-82: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winsuggestion: Add a
graph_ctxstructured-partition test.places/cute_parallel_for.cuexercisesstream_ctx, while this PR adds a separate graph allocation branch inslice_graph_interface; graph-only failures remain untested. As per path instructions, cudax reviews must focus on correctness and tests.Source: Path instructions
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: c9639889-676a-41b9-b0d6-0f0fab484257
📒 Files selected for processing (20)
cudax/examples/places/thrust_device_data_place_allocator.cucudax/examples/stf/CMakeLists.txtcudax/examples/stf/fdtd_mgpu.cucudax/examples/stf/partitioned_axpy.cucudax/include/cuda/experimental/__places/cute_partition.cuhcudax/include/cuda/experimental/__places/data_place_interface.cuhcudax/include/cuda/experimental/__places/localized_array.cuhcudax/include/cuda/experimental/__places/places.cuhcudax/include/cuda/experimental/__stf/graph/interfaces/slice.cuhcudax/include/cuda/experimental/__stf/internal/backend_ctx.cuhcudax/include/cuda/experimental/__stf/internal/parallel_for_scope.cuhcudax/include/cuda/experimental/__stf/internal/stf_places_extended_exports.cuhcudax/include/cuda/experimental/__stf/stream/interfaces/slice.cuhcudax/include/cuda/experimental/__stf/utility/dimensions.cuhcudax/include/cuda/experimental/places.cuhcudax/test/places/CMakeLists.txtcudax/test/places/placement.cucudax/test/stf/CMakeLists.txtcudax/test/stf/places/cute_parallel_for.cudocs/cudax/places.rst
🚧 Files skipped from review as they are similar to previous changes (18)
- cudax/include/cuda/experimental/places.cuh
- cudax/include/cuda/experimental/__stf/internal/backend_ctx.cuh
- cudax/examples/stf/CMakeLists.txt
- cudax/include/cuda/experimental/__places/data_place_interface.cuh
- cudax/include/cuda/experimental/__stf/internal/stf_places_extended_exports.cuh
- docs/cudax/places.rst
- cudax/examples/places/thrust_device_data_place_allocator.cu
- cudax/test/places/CMakeLists.txt
- cudax/examples/stf/fdtd_mgpu.cu
- cudax/include/cuda/experimental/__stf/utility/dimensions.cuh
- cudax/test/places/placement.cu
- cudax/examples/stf/partitioned_axpy.cu
- cudax/test/stf/places/cute_parallel_for.cu
- cudax/include/cuda/experimental/__stf/stream/interfaces/slice.cuh
- cudax/include/cuda/experimental/__places/places.cuh
- cudax/include/cuda/experimental/__stf/internal/parallel_for_scope.cuh
- cudax/include/cuda/experimental/__places/localized_array.cuh
- cudax/include/cuda/experimental/__places/cute_partition.cuh
2 tasks
caugonnet
force-pushed
the
stf-cute-parallel-for
branch
from
July 16, 2026 09:55
6a6fb2e to
7369680
Compare
The parallel_for machinery was generic over partitioners in shape but assumed statelessness: the instance received by the public entry point was dropped, apply() was invoked as a static, and the composite data place was built by default-constructing a fresh partitioner. Keep the instance instead ([[no_unique_address]], with an empty stand-in for the partitioner-less case), invoke apply through it (static-through- instance keeps the classic partitioners bit-identical), and construct the composite place from it - via the fn-pointer path when the partitioner has a static get_executor, or from the stateful object otherwise. cute_partition gains the partitioner-contract apply(): it derives the per-place sub-shape (cute_sub_shape - trivially copyable, iterated by the existing parallel_for kernels through size()/index_to_coords()). Restricted to exact covers for now: the shape interface has no notion of skipping the phantom coordinates padding would introduce, so uneven extents are rejected with an exception rather than silently computing on them. The slice data interfaces route composite places backed by a structured partition to a direct localized_array allocation (the stateful owner function cannot use the fn-pointer keyed recycling cache; recycling such arrays is a recorded follow-up - they are parked in the cache on deallocation and released at finalization). is_composite() becomes a virtual on data_place_interface (previously typeid-exact, which excluded the cute composite). The user-facing entry point is unchanged: the same ctx.parallel_for(partitioner, place, shape, deps...) now accepts both classic partitioners and cute_partition instances, and one partition object decides both the kernel decomposition and the data placement. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Covers 1-D blocked, the motivating 3-D dimension-1 split, the classic stateless partitioners passing through the same entry unchanged, and the exact-cover rejection of uneven extents at the apply() contract level (a partitioner throwing mid-task is not recoverable by design, so the runtime path is not exercised for the failure case). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Extend the cute_partition parallel_for path with the predication idiom (iterate each place's own coordinates, skip those outside the region) instead of restructuring iteration: - cute_sub_shape carries per-dimension [lo, hi) bounds and a contains() predicate; the parallel_for loops (device kernels and the host callback) skip coordinates outside it when the shape provides one. - apply() no longer rejects uneven extents: padding phantoms are excluded by the predicate, so odd sizes work end to end. - apply(box, place): the box is a region within the coordinate space of the tensor the partition was built for (the partition remains the authority on the extents, validated by containment). Each place iterates its owned coordinates restricted to the box, keeping the compute decomposition aligned with data ownership - unlike scale-free partitioners, which split the box itself and misalign with ownership at the region edges. Adapt fdtd_mgpu as the demonstration: one make_partition call decides which dimension splits for all ten tasks (init over the full shape, updates over interior boxes, the source point), and the same object places the data. The test gains uneven-extent and interior-region coverage, boundary values verified untouched. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Thin regions (boundary faces) invert the predication density trade-off, and per-place interval intersection is not worth its structural cost (apply works policy-agnostically on leaves; intersection would need per-dimension ownership knowledge). The supported pattern is simpler: iterate the face with a classic scale-free partitioner (tight, no discarded lanes) while explicit deps keep placement on the cute composite. This relies on separately constructed cute composites comparing equal (single instance identity across tasks), which the test now guards. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The compute-side companion of the structured-partitions section: the unchanged parallel_for entry point, regions (boxes) within the partition's tensor, the predication cost model (dense regions yes, thin regions no), the boundary-update menu, and the fdtd_mgpu demonstration. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The owner lambda called data_dims.index_to_pos() directly, bypassing the dimensions == 0 guard the shared delinearize helper provides for scalar slices. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Cedric AUGONNET <caugonnet@nvidia.com>
caugonnet
force-pushed
the
stf-cute-parallel-for
branch
from
July 16, 2026 17:04
7369680 to
6a0ecf3
Compare
Reuse CuTe-backed VMM allocations and their prerequisites through a value-keyed cache while preserving the classic partitioner cache path.
Describe value-defined CuTe ownership precisely and centralize the composite-place downcast used by cache operations.
Rely on the STF places bridge instead of repeating redundant namespace imports in the structured-partition coverage.
Reject rank and allocation extent mismatches, avoid dense point-source predication, and cover the graph backend.
Keep rank and leaf counts available to kernel code while erasing partitions only at runtime boundaries, reducing launch state without sacrificing cache or interop support.
Contributor
Author
|
@coderabbitai review |
Contributor
✅ Action performedReview finished.
|
Keep the internal cache protocol on cuda::std types and verify that structured allocations reject equal-size tensors with mismatched shapes.
Contributor
Author
|
/ok to test c9dbab9 |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
cudax/include/cuda/experimental/__places/cute_partition.cuh (1)
1165-1181: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winimportant: guard the partition padding arithmetic against overflow.
e.block * nplacescan wrap to zero, making the following division undefined; the ceil additions and final multiplication can also wrap. Use checked multiplication and overflow-safe ceiling division.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d0eee0b2-3463-4d4c-83b4-71420a595078
📒 Files selected for processing (18)
cudax/examples/stf/custom_data_interface.cucudax/examples/stf/fdtd_mgpu.cucudax/examples/stf/partitioned_axpy.cucudax/include/cuda/experimental/__places/cute_partition.cuhcudax/include/cuda/experimental/__places/data_place_interface.cuhcudax/include/cuda/experimental/__places/places.cuhcudax/include/cuda/experimental/__stf/graph/interfaces/slice.cuhcudax/include/cuda/experimental/__stf/internal/backend_ctx.cuhcudax/include/cuda/experimental/__stf/internal/parallel_for_scope.cuhcudax/include/cuda/experimental/__stf/internal/slice.cuhcudax/include/cuda/experimental/__stf/internal/stf_places_extended_exports.cuhcudax/include/cuda/experimental/__stf/localization/composite_slice.cuhcudax/include/cuda/experimental/__stf/stream/interfaces/slice.cuhcudax/include/cuda/experimental/__stf/utility/dimensions.cuhcudax/test/places/placement.cucudax/test/stf/CMakeLists.txtcudax/test/stf/places/cute_parallel_for.cudocs/cudax/places.rst
🚧 Files skipped from review as they are similar to previous changes (7)
- cudax/include/cuda/experimental/__stf/internal/stf_places_extended_exports.cuh
- cudax/test/places/placement.cu
- cudax/include/cuda/experimental/__stf/internal/backend_ctx.cuh
- cudax/examples/stf/fdtd_mgpu.cu
- cudax/test/stf/CMakeLists.txt
- cudax/include/cuda/experimental/__stf/internal/parallel_for_scope.cuh
- cudax/examples/stf/partitioned_axpy.cu
This comment has been minimized.
This comment has been minimized.
Use host standard traits for the coordinate-predicate detector because NVCC 12.0 fails to discard the cuda::std::void_t specialization for shapes without contains().
Reject reversed iteration boxes before dispatch and ensure cache/equality coverage uses independently constructed equivalent partitions.
Contributor
Author
|
/ok to test e2f011d |
Use overload-based shape predication so NVCC 12 never instantiates a missing contains() member, including inside device lambdas.
Contributor
Author
|
/ok to test 7bd0aa1 |
Use std::pair for the host-only cache API to avoid CUDA 12 compiling event_list's move path for device code.
Contributor
Author
|
/ok to test 1d9cd7c |
Contributor
🥳 CI Workflow Results🟩 Finished in 1h 10m: Pass: 100%/57 | Total: 1d 04h | Max: 1h 10m | Hits: 13%/187274See results here. |
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
2 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
Built on #9804, now merged. This PR contains only the six compute-integration follow-up commits.
This completes the "one object drives compute and data" arc of #9804: the same user-facing
ctx.parallel_for(partitioner, place, shape, deps…)now accepts acute_partitioninstance, which decides both the per-place kernel decomposition and (through the task's affine composite data place) the physical placement of the data those kernels touch.apply()was invoked as a static, and the composite data place was default-constructed. The instance is now kept ([[no_unique_address]], an empty stand-in for the partitioner-less case) and used at both sites. The classic partitioners compile through the identical path (static-through-instance invocation) — no behavior change for any existingparallel_for.cute_partition::apply()returns a trivially-copyablecute_sub_shapeenumerating the place's own coordinates, with per-dimension[lo, hi)bounds consumed by acontains()predicate that theparallel_forloops honor when a shape provides one (if constexpron a trait — codegen for existing shapes is unchanged). This handles padding phantoms (uneven extents work end to end) and interior regions (apply(box, place): the box is a region within the tensor the partition was built for, containment-validated; each place iterates its owned coordinates restricted to the box, keeping compute aligned with data ownership). SASS-verified: kernels contain zero calls — decode, predicate and body are one flat inlined stream, the predicate being a handful ofISETPs against an iteration dominated by the pre-existing decode divmods.localized_arrayallocation (a stateful owner cannot use the fn-pointer-keyed recycling cache; such arrays park in the cache at deallocation and free at finalization — recycling them is a recorded follow-up).is_composite()becomes a virtual ondata_place_interface(was typeid-exact).fdtd_mgpuadapted as the demonstration: onemake_partitioncall decides which dimension splits for all ten tasks — init over the full shape, updates over interior boxes, the source point — and the same object places the data. Changing the split dimension is editing one spec entry, which the classicblocked_partition(last-dimension-only) cannot express.Tests:
cudax/test/stf/places/cute_parallel_for.cu— 1-D blocked, 3-D dimension-1 split, classic partitioner passthrough, uneven extents through the runtime, interior region with untouched boundary, and the mixed boundary pattern.Recorded follow-ups: run-structured decode in
cute_sub_shape(amortizes the per-element divmods — the dominant arithmetic per the SASS), recycling partition-backed arrays in the composite cache,optionally_staticleaf/extent refinement.Checklist
🤖 Generated with Claude Code