From 09c5c3e57a7bf7d53b06ef3b721187bdd6ec1789 Mon Sep 17 00:00:00 2001 From: Cedric AUGONNET Date: Sat, 15 Aug 2026 10:58:50 +0000 Subject: [PATCH 01/15] [STF] Add bundles: grouped logical data as a single task dependency A bundle ties several logical data together behind one object (a CSR matrix's three arrays, a graph's topology) submitted as ONE dependency: the construct expands it into ordinary per-field dependencies, and the user lambda receives one tuple of views per bundle. Every field remains a first-class logical data (a bundle owns no data and no tracking state), so bundle-level and bare-handle dependencies interoperate by meeting at the same logical data. Fields declared 'constant' carry a read-only ceiling: whole-bundle modes distribute per field as the strongest admitted mode (rw() on a bundle clamps constant fields to read), and constant fields are const-qualified in every view via the existing readonly_type_of mapping. write() also clamps constant fields to read (a writer needs the structure fetched to interpret what it writes). Mechanism: header-only adapter layer. bundle_dep expands at the context entry points (context + backend_ctx overloads, SFINAE-gated on any_bundle_dep_v so existing overloads are untouched); a named host-device functor regroups the flat views into per-bundle tuples in front of the user function, with its call operator constrained to the wrapped function's applicability so convention probing keeps working; the extended-lambda classification in parallel_for/launch looks through the adapter at the wrapped lambda (identity for everything else). Examples: cg_csr.cu's csr_matrix becomes a bundle (SpMV: 5 deps -> 3, structure constant); pagerank.cu groups the CSR topology as a constant bundle next to its reduce dependency. Test: interface/bundle.cu covers adopt/create constructors, mode distribution and const views (compile-time), mixed bundle+bare-leaf use, task/parallel_for/ host_launch, stream and graph backends. Verified on GB300 sm_103 / CUDA 13.4: new test + both examples build and run under c++17 and c++20; launch/reduce/token regression files recompile clean. Co-Authored-By: Claude Fable 5 --- .../examples/stf/graph_algorithms/pagerank.cu | 21 +- cudax/examples/stf/linear_algebra/cg_csr.cu | 55 +- .../__stf/internal/backend_ctx.cuh | 82 ++- .../experimental/__stf/internal/bundle.cuh | 468 ++++++++++++++++++ .../experimental/__stf/internal/context.cuh | 104 +++- .../experimental/__stf/internal/launch.cuh | 11 +- .../__stf/internal/parallel_for_scope.cuh | 6 +- cudax/test/stf/CMakeLists.txt | 1 + cudax/test/stf/interface/bundle.cu | 122 +++++ 9 files changed, 826 insertions(+), 44 deletions(-) create mode 100644 cudax/include/cuda/experimental/__stf/internal/bundle.cuh create mode 100644 cudax/test/stf/interface/bundle.cu diff --git a/cudax/examples/stf/graph_algorithms/pagerank.cu b/cudax/examples/stf/graph_algorithms/pagerank.cu index c39cdd686da7..b724b30b2539 100644 --- a/cudax/examples/stf/graph_algorithms/pagerank.cu +++ b/cudax/examples/stf/graph_algorithms/pagerank.cu @@ -67,8 +67,14 @@ int main() std::vector page_rank(num_vertices, init_rank); std::vector new_page_rank(num_vertices); - auto loffsets = ctx.logical_data(&offsets[0], offsets.size()); - auto lnonzeros = ctx.logical_data(&nonzeros[0], nonzeros.size()); + auto loffsets = ctx.logical_data(&offsets[0], offsets.size()); + auto lnonzeros = ctx.logical_data(&nonzeros[0], nonzeros.size()); + + // The CSR graph topology is one object: a bundle of the two constant arrays. + // Tasks depend on `graph` as a whole and receive one tuple of (const) views. + bundle, constant>, field, constant>> graph(loffsets, lnonzeros); + graph.set_symbol("graph"); + auto lpage_rank = ctx.logical_data(&page_rank[0], page_rank.size()); auto lnew_page_rank = ctx.logical_data(&new_page_rank[0], new_page_rank.size()); auto lmax_diff = ctx.logical_data(shape_of>()); @@ -77,14 +83,9 @@ int main() { // Calculate Current Iteration PageRank ctx.parallel_for( - box(num_vertices), - loffsets.read(), - lnonzeros.read(), - lpage_rank.rw(), - lnew_page_rank.rw(), - lmax_diff.reduce(reducer::maxval{})) - ->*[init_rank] __device__( - size_t idx, auto loffsets, auto lnonzeros, auto lpage_rank, auto lnew_page_rank, auto& max_diff) { + box(num_vertices), graph.read(), lpage_rank.rw(), lnew_page_rank.rw(), lmax_diff.reduce(reducer::maxval{})) + ->*[init_rank] __device__(size_t idx, auto graph, auto lpage_rank, auto lnew_page_rank, auto& max_diff) { + auto& [loffsets, lnonzeros] = graph; calculating_pagerank(idx, loffsets, lnonzeros, lpage_rank, lnew_page_rank, init_rank); max_diff = ::std::max(max_diff, lnew_page_rank[idx] - lpage_rank[idx]); }; diff --git a/cudax/examples/stf/linear_algebra/cg_csr.cu b/cudax/examples/stf/linear_algebra/cg_csr.cu index e88e6b578c64..68bf1875d4c4 100644 --- a/cudax/examples/stf/linear_algebra/cg_csr.cu +++ b/cudax/examples/stf/linear_algebra/cg_csr.cu @@ -21,20 +21,30 @@ using vector_t = logical_data>; using scalar_t = logical_data>; using context_t = context; -struct csr_matrix +/* A CSR matrix is one object made of three arrays: a bundle of the values + * and of the (constant) structure. Tasks depend on the whole matrix with a + * single argument, and every field remains an ordinary logical data. */ +struct csr_matrix : bundle>, field, constant>, field, constant>> { csr_matrix( context_t& ctx, size_t num_rows, size_t num_nonzeros, double* values, size_t* row_offsets, size_t* column_indices) + : bundle(ctx.logical_data(make_slice(values, num_nonzeros)), + ctx.logical_data(make_slice(column_indices, num_nonzeros)), + ctx.logical_data(make_slice(row_offsets, num_rows + 1))) + {} + + auto& vals() { - val_handle = ctx.logical_data(make_slice(values, num_nonzeros)); - col_handle = ctx.logical_data(make_slice(column_indices, num_nonzeros)); - row_handle = ctx.logical_data(make_slice(row_offsets, num_rows + 1)); + return get_field<0>(); + } + auto& colind() + { + return get_field<1>(); + } + auto& rowptr() + { + return get_field<2>(); } - - /* Description of the CSR */ - mutable logical_data> val_handle; - mutable logical_data> row_handle; - mutable logical_data> col_handle; }; // Note that a and b might be the same logical data @@ -48,19 +58,20 @@ void DOT(context_t& ctx, vector_t& a, vector_t& b, scalar_t& res) void SPMV(context_t& ctx, csr_matrix& a, vector_t& x, vector_t& y) { - ctx.parallel_for(y.shape(), a.val_handle.read(), a.col_handle.read(), a.row_handle.read(), x.read(), y.write()) - ->*[] _CCCL_DEVICE(size_t row, auto da_val, auto da_col, auto da_row, auto dx, auto dy) { - int row_start = da_row(row); - int row_end = da_row(row + 1); - - double sum = 0.0; - for (int elt = row_start; elt < row_end; elt++) - { - sum += da_val(elt) * dx(da_col(elt)); - } - - dy(row) = sum; - }; + ctx.parallel_for(y.shape(), a.read(), x.read(), y.write())->*[] _CCCL_DEVICE(size_t row, auto da, auto dx, auto dy) { + auto& [da_val, da_col, da_row] = da; + + int row_start = da_row(row); + int row_end = da_row(row + 1); + + double sum = 0.0; + for (int elt = row_start; elt < row_end; elt++) + { + sum += da_val(elt) * dx(da_col(elt)); + } + + dy(row) = sum; + }; } /* genTridiag: generate a random tridiagonal symmetric matrix diff --git a/cudax/include/cuda/experimental/__stf/internal/backend_ctx.cuh b/cudax/include/cuda/experimental/__stf/internal/backend_ctx.cuh index f194ba6a1953..0e501e9c2304 100644 --- a/cudax/include/cuda/experimental/__stf/internal/backend_ctx.cuh +++ b/cudax/include/cuda/experimental/__stf/internal/backend_ctx.cuh @@ -33,6 +33,7 @@ #include #include #include +#include // bundle-aware dependency overloads #include #include // backend_ctx::launch() uses execution_policy #include @@ -1151,7 +1152,8 @@ public: template >> + typename = ::cuda::std::enable_if_t<::cuda::std::is_base_of_v + && !reserved::any_bundle_dep_v>> auto parallel_for(exec_place_t e_place, S shape, Deps... deps) { if constexpr (::cuda::std::is_integral_v) @@ -1169,7 +1171,8 @@ public: typename exec_place_t, typename S, typename... Deps, - typename = ::cuda::std::enable_if_t<::cuda::std::is_base_of_v>> + typename = ::cuda::std::enable_if_t<::cuda::std::is_base_of_v + && !reserved::any_bundle_dep_v>> auto parallel_for(partitioner_t p, exec_place_t e_place, S shape, Deps... deps) { if constexpr (::cuda::std::is_integral_v) @@ -1183,12 +1186,85 @@ public: } } - template + template >> auto parallel_for(S shape, Deps... deps) { return parallel_for(self().default_exec_place(), mv(shape), mv(deps)...); } + /* + * Bundle-aware overloads: when a dependency list contains bundle + * dependencies, expand them into their per-field dependencies, delegate to + * the ordinary construct, and wrap the resulting scope so the user function + * receives one tuple of views per bundle. + */ + template >> + auto task(exec_place e_place, Args... args) + { + return reserved::make_bundle_scope( + [&](auto... flat) { + return self().task(mv(e_place), mv(flat)...); + }, + mv(args)...); + } + + template >> + auto task(Args... args) + { + return task(self().default_exec_place(), mv(args)...); + } + + template < + typename exec_place_t, + typename S, + typename... Args, + ::cuda::std::enable_if_t<::cuda::std::is_base_of_v && reserved::any_bundle_dep_v, + int> = 0> + auto parallel_for(exec_place_t e_place, S shape, Args... args) + { + return reserved::make_bundle_scope( + [&](auto... flat) { + return this->parallel_for(mv(e_place), mv(shape), mv(flat)...); + }, + mv(args)...); + } + + template , int> = 0> + auto parallel_for(S shape, Args... args) + { + return parallel_for(self().default_exec_place(), mv(shape), mv(args)...); + } + + template >> + auto host_launch(Args... args) + { + return reserved::make_bundle_scope( + [&](auto... flat) { + return this->host_launch(mv(flat)...); + }, + mv(args)...); + } + + template >> + auto cuda_kernel(Args... args) + { + return reserved::make_bundle_scope( + [&](auto... flat) { + return this->cuda_kernel(mv(flat)...); + }, + mv(args)...); + } + + template >> + auto cuda_kernel_chain(Args... args) + { + return reserved::make_bundle_scope( + [&](auto... flat) { + return this->cuda_kernel_chain(mv(flat)...); + }, + mv(args)...); + } + private: Engine& self() { diff --git a/cudax/include/cuda/experimental/__stf/internal/bundle.cuh b/cudax/include/cuda/experimental/__stf/internal/bundle.cuh new file mode 100644 index 000000000000..f2b8cf1d6bc4 --- /dev/null +++ b/cudax/include/cuda/experimental/__stf/internal/bundle.cuh @@ -0,0 +1,468 @@ +//===----------------------------------------------------------------------===// +// +// Part of CUDASTF in CUDA C++ Core Libraries, +// under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +// +//===----------------------------------------------------------------------===// + +/** + * @file + * + * @brief Bundles: non-owning groups of logical data usable as a single task dependency + * + * A `bundle` ties several logical data together behind one object (e.g. the + * three arrays of a CSR matrix) so that tasks can depend on the whole group + * with a single argument, while every constituent ("field") remains an + * ordinary logical data usable on its own. A bundle owns no data and no + * dependency-tracking state: it merely holds handle copies, and a bundle + * dependency expands into one ordinary dependency per field. The lambda of a + * task (or parallel_for, ...) receives one tuple of per-field views per + * bundle dependency instead of one view per field. + * + * Fields declared `constant` have a read-only ceiling: whole-bundle access + * modes distribute to each field as the strongest mode the field admits + * (`rw()` on a bundle with a constant field passes that field's view as + * read-only), and their views are const-qualified in every spelling. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include + +#include +#include + +#include +#include +#include + +namespace cuda::experimental::stf +{ +// Bundles only ever name these inside templates; the definitions come from +// logical_data.cuh / data_interface.cuh at instantiation time. Keeping this +// header lightweight lets backend_ctx.cuh include it without cycles. +template +class logical_data; + +template +class shape_of; + +/** + * @brief Trait tag marking a bundle field as constant (read-only ceiling). + * + * A constant field only ever admits read access: whole-bundle modes clamp to + * `read` for it, and its view is const-qualified in every task. + */ +struct constant +{}; + +/** + * @brief Describes one field of a bundle: the instance type and optional traits. + * + * @tparam T the instance type of the underlying logical data (e.g. `slice`) + * @tparam Traits optional traits (`constant_t`) + */ +template +struct field +{ + using type = T; + static constexpr bool is_constant = (::cuda::std::is_same_v || ... || false); +}; + +namespace reserved +{ +template +struct is_field : ::cuda::std::false_type +{}; + +template +struct is_field> : ::cuda::std::true_type +{}; +} // end namespace reserved + +/** + * @brief A group of ordinary task dependencies submitted as a single argument. + * + * Produced by `bundle::read()/rw()/write()`; consumed by the context task + * constructs, which expand it into its per-field dependencies and regroup the + * corresponding views into a single tuple argument for the user lambda. + */ +template +class bundle_dep +{ +public: + static constexpr size_t arity = sizeof...(LeafDeps); + + explicit bundle_dep(LeafDeps... d) + : deps(mv(d)...) + {} + + ::std::tuple deps; +}; + +namespace reserved +{ +template +struct is_bundle_dep : ::cuda::std::false_type +{}; + +template +struct is_bundle_dep> : ::cuda::std::true_type +{}; + +template +inline constexpr bool is_bundle_dep_v = is_bundle_dep<::cuda::std::remove_cvref_t>::value; + +//! True when at least one argument of a dependency list is a bundle_dep +template +inline constexpr bool any_bundle_dep_v = (is_bundle_dep_v || ... || false); + +//! Number of flat dependencies one submitted argument expands to +template +struct slot_arity : ::cuda::std::integral_constant +{}; + +template +struct slot_arity> : ::cuda::std::integral_constant +{}; + +//! Expand one submitted dependency argument into a tuple of flat dependencies +template +auto as_dep_tuple(T t) +{ + if constexpr (is_bundle_dep_v) + { + return mv(t.deps); + } + else + { + return ::std::make_tuple(mv(t)); + } +} +} // end namespace reserved + +/** + * @brief A non-owning group of logical data described by a list of `field`s. + * + * Bundles introduce no new ownership domain: they hold ordinary (refcounted) + * logical data handles. Every field remains a first-class logical data + * retrievable with `get_field()`, so bundle-level and per-field task + * dependencies interoperate (they meet at the same logical data). + * + * @tparam Fields `field` descriptors, in canonical order + */ +template +class bundle +{ + static_assert((reserved::is_field::value && ...), "bundle<...> parameters must be field<...> descriptors"); + +public: + static constexpr size_t n_fields = sizeof...(Fields); + + template + using field_at = ::std::tuple_element_t>; + + /** @brief Adopt existing logical data as the bundle's fields (no context needed) */ + explicit bundle(logical_data... h) + : handles(mv(h)...) + {} + + /** @brief Create fresh logical data from shapes, then behave like an adopting bundle */ + template + bundle(ctx_t& ctx, shape_of... shapes) + : handles(ctx.logical_data(mv(shapes))...) + {} + + /** @brief Access field `I` as its plain, first-class logical data */ + template + auto& get_field() + { + return ::std::get(handles); + } + + template + const auto& get_field() const + { + return ::std::get(handles); + } + + /** @brief Depend on every field with read access */ + auto read() + { + return dep_impl(::std::index_sequence_for()); + } + + /** + * @brief Depend on the bundle with read-write access. + * + * Distributes per field as the strongest admitted mode: mutable fields get + * `rw`, constant fields clamp to `read` (their views stay const-qualified). + */ + auto rw() + { + return dep_impl(::std::index_sequence_for()); + } + + /** + * @brief Depend on the bundle with write access. + * + * Mutable fields get `write` (previous content is discarded, not fetched); + * constant fields clamp to `read` — their content is still fetched, since a + * writer typically needs the constant fields to interpret what it writes. + */ + auto write() + { + return dep_impl(::std::index_sequence_for()); + } + + /** @brief Name unnamed fields "prefix." (never clobbers an existing symbol) */ + bundle& set_symbol(const ::std::string& prefix) + { + set_symbol_impl(prefix, ::std::index_sequence_for()); + return *this; + } + +private: + template + auto leaf_dep() + { + auto& ld = ::std::get(handles); + if constexpr (M == access_mode::read || field_at::is_constant) + { + return ld.read(); + } + else if constexpr (M == access_mode::write) + { + return ld.write(); + } + else + { + return ld.rw(); + } + } + + template + auto dep_impl(::std::index_sequence) + { + return bundle_deptemplate leaf_dep())...>(this->template leaf_dep()...); + } + + template + void set_symbol_impl(const ::std::string& prefix, ::std::index_sequence) + { + ((::std::get(handles).get_symbol().empty() + ? void(::std::get(handles).set_symbol(prefix + "." + ::std::to_string(Is))) + : void()), + ...); + } + + ::std::tuple...> handles; +}; + +namespace reserved +{ +/** + * @brief Callable adapter regrouping the flat per-dependency views of an + * expanded dependency list back into one tuple argument per bundle. + * + * The wrapped construct passes `(leading..., view0, view1, ...)` where the + * trailing `total_flat` arguments are the flat per-dependency views in + * submission order and the leading arguments are construct-specific (a + * stream, shape coordinates, a thread hierarchy, ...). Slots of arity 1 are + * forwarded untouched (preserving references); slots of arity k become a + * `::cuda::std::tuple` of the k views, by value. + * + * This is a named functor (not a lambda) so it can wrap extended device + * lambdas, and its call operator is SFINAE-constrained to the wrapped + * function's own applicability so that constructs probing several calling + * conventions keep working. + */ +template +struct bundle_arg_adapter; + +template +struct bundle_arg_adapter> +{ + Fun f; + + static constexpr size_t n_slots = sizeof...(Arities); + static constexpr size_t total_flat = (Arities + ... + 0); + + static constexpr size_t arity_of(size_t slot) + { + constexpr size_t a[] = {Arities...}; + return a[slot]; + } + + static constexpr size_t offset_of(size_t slot) + { + constexpr size_t a[] = {Arities...}; + size_t o = 0; + for (size_t i = 0; i < slot; ++i) + { + o += a[i]; + } + return o; + } + + // Group the slot starting at flat position Offset (relative to the first + // flat view) out of the full forwarded argument tuple. NLead is the number + // of leading (non-dependency) arguments. + _CCCL_EXEC_CHECK_DISABLE + template + _CCCL_HOST_DEVICE static decltype(auto) group_slot(Tup& tup) + { + constexpr size_t off = NLead + offset_of(Slot); + if constexpr (arity_of(Slot) == 1) + { + // Forward the argument as-is (a reduction accumulator must remain a reference) + return static_cast<::cuda::std::tuple_element_t>&&>( + ::cuda::std::get(tup)); + } + else + { + return make_view_tuple(tup, ::std::make_index_sequence()); + } + } + + _CCCL_EXEC_CHECK_DISABLE + template + _CCCL_HOST_DEVICE static auto make_view_tuple(Tup& tup, ::std::index_sequence) + { + return ::cuda::std::tuple< + ::cuda::std::remove_cvref_t<::cuda::std::tuple_element_t>>...>( + ::cuda::std::get(tup)...); + } + + // Type-level mirror of group_slot, used to constrain operator() + template + using group_slot_t = decltype(group_slot(::cuda::std::declval())); + + template + static constexpr bool invocable_impl(::std::index_sequence, ::std::index_sequence) + { + return ::cuda::std:: + is_invocable_v&&..., group_slot_t...>; + } + + template + static constexpr bool applicable() + { + if constexpr (sizeof...(Args) < total_flat) + { + return false; + } + else + { + constexpr size_t n_lead = sizeof...(Args) - total_flat; + return invocable_impl<::cuda::std::tuple, n_lead>( + ::std::make_index_sequence(), ::std::make_index_sequence()); + } + } + + _CCCL_EXEC_CHECK_DISABLE + template + _CCCL_HOST_DEVICE decltype(auto) call(Tup&& tup, ::std::index_sequence, ::std::index_sequence) + { + constexpr size_t n_lead = sizeof...(LeadIs); + return f(::cuda::std::get(mv(tup))..., group_slot(tup)...); + } + + _CCCL_EXEC_CHECK_DISABLE + template ()>> + _CCCL_HOST_DEVICE decltype(auto) operator()(Args&&... args) + { + constexpr size_t n_lead = sizeof...(Args) - total_flat; + return call(::cuda::std::forward_as_tuple(::cuda::std::forward(args)...), + ::std::make_index_sequence(), + ::std::make_index_sequence()); + } +}; + +//! Classification helper: the executable "kind" of an adapter (extended +//! device lambda, host-device lambda, plain host callable) is that of the +//! user function it wraps. Identity for every other callable. +template +struct bundle_inner_fun +{ + using type = F; +}; + +template +struct bundle_inner_fun> +{ + using type = typename bundle_inner_fun::type; +}; + +template +using bundle_inner_fun_t = typename bundle_inner_fun<::cuda::std::remove_cvref_t>::type; + +/** + * @brief Wraps a task-like construct scope so that `operator->*` regroups + * bundle views into single tuple arguments before invoking the user function. + */ +template +class bundle_scope +{ +public: + explicit bundle_scope(Scope s) + : inner(mv(s)) + {} + + bundle_scope& set_symbol(::std::string s) & + { + inner.set_symbol(mv(s)); + return *this; + } + + bundle_scope&& set_symbol(::std::string s) && + { + inner.set_symbol(mv(s)); + return mv(*this); + } + + Scope& get_scope() + { + return inner; + } + + template + decltype(auto) operator->*(Fun&& f) + { + return mv(inner)->*bundle_arg_adapter<::cuda::std::remove_cvref_t, AritySeq>{::cuda::std::forward(f)}; + } + +private: + Scope inner; +}; + +/** + * @brief Expand a dependency list that contains bundle_dep arguments into + * flat dependencies, invoke `make_inner` on them, and wrap the resulting + * scope for view regrouping. + */ +template +auto make_bundle_scope(MakeInner&& make_inner, Args... args) +{ + using arities = ::std::index_sequence::value...>; + auto flat = ::std::tuple_cat(as_dep_tuple(mv(args))...); + auto inner = ::std::apply(::cuda::std::forward(make_inner), mv(flat)); + return bundle_scope(mv(inner)); +} +} // end namespace reserved +} // end namespace cuda::experimental::stf diff --git a/cudax/include/cuda/experimental/__stf/internal/context.cuh b/cudax/include/cuda/experimental/__stf/internal/context.cuh index 204842f04f21..a2b0e0b1045d 100644 --- a/cudax/include/cuda/experimental/__stf/internal/context.cuh +++ b/cudax/include/cuda/experimental/__stf/internal/context.cuh @@ -27,6 +27,7 @@ #include #include #include +#include #include #include #include @@ -673,7 +674,8 @@ public: template >> + typename = ::cuda::std::enable_if_t<::cuda::std::is_base_of_v + && !reserved::any_bundle_dep_v>> auto parallel_for(exec_place_t e_place, S shape, Deps... deps) { if constexpr (::cuda::std::is_integral_v) @@ -695,7 +697,8 @@ public: typename exec_place_t, typename S, typename... Deps, - typename = ::cuda::std::enable_if_t<::cuda::std::is_base_of_v>> + typename = ::cuda::std::enable_if_t<::cuda::std::is_base_of_v + && !reserved::any_bundle_dep_v>> auto parallel_for(partitioner_t p, exec_place_t e_place, S shape, Deps... deps) { EXPECT(payload.index() != ::cuda::std::variant_npos, "Context is not initialized."); @@ -706,7 +709,7 @@ public: }; } - template + template >> auto parallel_for(S shape, Deps... deps) { return parallel_for(default_exec_place(), mv(shape), mv(deps)...); @@ -805,6 +808,101 @@ public: } #endif // !defined(CUDASTF_DISABLE_CODE_GENERATION) && _CCCL_CUDA_COMPILATION() + /* + * Bundle-aware overloads: when a dependency list contains bundle + * dependencies, expand them into their per-field dependencies, delegate to + * the ordinary construct, and wrap the resulting scope so the user function + * receives one tuple of views per bundle. + */ + template >> + auto task(exec_place e_place, Args... args) + { + return reserved::make_bundle_scope( + [&](auto... flat) { + return this->task(mv(e_place), mv(flat)...); + }, + mv(args)...); + } + + template >> + auto task(Args... args) + { + return task(default_exec_place(), mv(args)...); + } + +#if !defined(CUDASTF_DISABLE_CODE_GENERATION) && _CCCL_CUDA_COMPILATION() + template < + typename exec_place_t, + typename S, + typename... Args, + ::cuda::std::enable_if_t<::cuda::std::is_base_of_v && reserved::any_bundle_dep_v, + int> = 0> + auto parallel_for(exec_place_t e_place, S shape, Args... args) + { + return reserved::make_bundle_scope( + [&](auto... flat) { + return this->parallel_for(mv(e_place), mv(shape), mv(flat)...); + }, + mv(args)...); + } + + template , int> = 0> + auto parallel_for(S shape, Args... args) + { + return parallel_for(default_exec_place(), mv(shape), mv(args)...); + } +#endif // !defined(CUDASTF_DISABLE_CODE_GENERATION) && _CCCL_CUDA_COMPILATION() + + template >> + auto host_launch(Args... args) + { + return reserved::make_bundle_scope( + [&](auto... flat) { + return this->host_launch(mv(flat)...); + }, + mv(args)...); + } + + template >> + auto cuda_kernel(Args... args) + { + return reserved::make_bundle_scope( + [&](auto... flat) { + return this->cuda_kernel(mv(flat)...); + }, + mv(args)...); + } + + template >> + auto cuda_kernel(exec_place e_place, Args... args) + { + return reserved::make_bundle_scope( + [&](auto... flat) { + return this->cuda_kernel(mv(e_place), mv(flat)...); + }, + mv(args)...); + } + + template >> + auto cuda_kernel_chain(Args... args) + { + return reserved::make_bundle_scope( + [&](auto... flat) { + return this->cuda_kernel_chain(mv(flat)...); + }, + mv(args)...); + } + + template >> + auto cuda_kernel_chain(exec_place e_place, Args... args) + { + return reserved::make_bundle_scope( + [&](auto... flat) { + return this->cuda_kernel_chain(mv(e_place), mv(flat)...); + }, + mv(args)...); + } + auto repeat(size_t count) { using result_t = unified_scope, reserved::repeat_scope>; diff --git a/cudax/include/cuda/experimental/__stf/internal/launch.cuh b/cudax/include/cuda/experimental/__stf/internal/launch.cuh index 7cb8b77e3eaa..1ebdc88e2c9d 100644 --- a/cudax/include/cuda/experimental/__stf/internal/launch.cuh +++ b/cudax/include/cuda/experimental/__stf/internal/launch.cuh @@ -25,6 +25,7 @@ #include #include +#include // bundle_inner_fun_t for lambda-kind classification #include // launch_impl() uses execution_policy #include #include @@ -217,8 +218,9 @@ public: # if _CCCL_CUDA_COMPILER(NVCC) // With nvcc, dedicated traits tell how a lambda can be executed. static constexpr bool is_extended_host_device_lambda_closure_type = - __nv_is_extended_host_device_lambda_closure_type(Fun), - is_extended_device_lambda_closure_type = __nv_is_extended_device_lambda_closure_type(Fun); + __nv_is_extended_host_device_lambda_closure_type(reserved::bundle_inner_fun_t), + is_extended_device_lambda_closure_type = + __nv_is_extended_device_lambda_closure_type(reserved::bundle_inner_fun_t); # else // ^^^ _CCCL_CUDA_COMPILER(NVCC) ^^^ / VVV !_CCCL_CUDA_COMPILER(NVCC) VVV // Only nvcc offers those traits. The claim below holds for nvc++, where every lambda can // indeed run on host and device. For clang-cuda it is provisional: a device-only lambda @@ -360,8 +362,9 @@ public: # if _CCCL_CUDA_COMPILER(NVCC) // With nvcc, dedicated traits tell how a lambda can be executed. static constexpr bool is_extended_host_device_lambda_closure_type = - __nv_is_extended_host_device_lambda_closure_type(Fun), - is_extended_device_lambda_closure_type = __nv_is_extended_device_lambda_closure_type(Fun); + __nv_is_extended_host_device_lambda_closure_type(reserved::bundle_inner_fun_t), + is_extended_device_lambda_closure_type = + __nv_is_extended_device_lambda_closure_type(reserved::bundle_inner_fun_t); # else // ^^^ _CCCL_CUDA_COMPILER(NVCC) ^^^ / VVV !_CCCL_CUDA_COMPILER(NVCC) VVV // Only nvcc offers those traits. The claim below holds for nvc++, where every lambda can // indeed run on host and device. For clang-cuda it is provisional: a device-only lambda diff --git a/cudax/include/cuda/experimental/__stf/internal/parallel_for_scope.cuh b/cudax/include/cuda/experimental/__stf/internal/parallel_for_scope.cuh index a5903173a90e..36021f63a1f7 100644 --- a/cudax/include/cuda/experimental/__stf/internal/parallel_for_scope.cuh +++ b/cudax/include/cuda/experimental/__stf/internal/parallel_for_scope.cuh @@ -33,6 +33,7 @@ #include #include // for null_partition +#include #include #include #include @@ -784,8 +785,9 @@ public: # if _CCCL_CUDA_COMPILER(NVCC) // With nvcc, dedicated traits tell how a lambda can be executed. static constexpr bool is_extended_host_device_lambda_closure_type = - __nv_is_extended_host_device_lambda_closure_type(Fun), - is_extended_device_lambda_closure_type = __nv_is_extended_device_lambda_closure_type(Fun); + __nv_is_extended_host_device_lambda_closure_type(reserved::bundle_inner_fun_t), + is_extended_device_lambda_closure_type = + __nv_is_extended_device_lambda_closure_type(reserved::bundle_inner_fun_t); # else // ^^^ _CCCL_CUDA_COMPILER(NVCC) ^^^ / vvv !_CCCL_CUDA_COMPILER(NVCC) // Only nvcc offers those traits. The claim below holds for nvc++, where every lambda can // indeed run on host and device. For clang-cuda it is provisional: a device-only lambda diff --git a/cudax/test/stf/CMakeLists.txt b/cudax/test/stf/CMakeLists.txt index 533ec10067d3..c61590b4868b 100644 --- a/cudax/test/stf/CMakeLists.txt +++ b/cudax/test/stf/CMakeLists.txt @@ -51,6 +51,7 @@ set( graph/graph_ctx_low_level.cu graph/static_graph_ctx.cu hashtable/test.cu + interface/bundle.cu interface/cuda_kernel_chain-add_deps.cu interface/cuda_kernel_chain-add_deps_low_level.cu interface/cuda_kernel_empty_args.cu diff --git a/cudax/test/stf/interface/bundle.cu b/cudax/test/stf/interface/bundle.cu new file mode 100644 index 000000000000..99e801690fd5 --- /dev/null +++ b/cudax/test/stf/interface/bundle.cu @@ -0,0 +1,122 @@ +//===----------------------------------------------------------------------===// +// +// Part of CUDASTF in CUDA C++ Core Libraries, +// under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +// +//===----------------------------------------------------------------------===// + +/** + * @file + * + * @brief Bundles: grouped logical data used as a single dependency + * + * Checks that a bundle dependency expands to per-field dependencies, that the + * user function receives one tuple of views per bundle, that constant fields + * are clamped to read access (const views) in every spelling, that fields + * remain usable as bare logical data concurrently with bundle use, and that + * all of this holds on both the stream and graph backends. + */ + +#include + +using namespace cuda::experimental::stf; + +using test_bundle = bundle>, field, constant>>; + +// The rw view must keep mutable fields mutable and const-qualify constant fields +template +__host__ __device__ void check_rw_view_types() +{ + static_assert(::cuda::std::is_same_v<::cuda::std::tuple_element_t<0, BundleView>, slice>, + "mutable field must stay mutable in an rw bundle view"); + static_assert(::cuda::std::is_same_v<::cuda::std::tuple_element_t<1, BundleView>, slice>, + "constant field must be const in an rw bundle view"); +} + +void run(context& ctx) +{ + const size_t N = 64; + + ::std::vector vals(N, 2.0); + ::std::vector idx(N); + for (size_t i = 0; i < N; i++) + { + idx[i] = static_cast(i); + } + + auto lv = ctx.logical_data(&vals[0], N); + auto li = ctx.logical_data(&idx[0], N); + auto lo = ctx.logical_data(shape_of>(N)); + + // Adopting constructor: no context needed, handles are shared + test_bundle B(lv, li); + B.set_symbol("B"); + + // Fields remain first-class logical data + static_assert(test_bundle::n_fields == 2); + + // parallel_for with an rw bundle dep: one tuple argument + ctx.parallel_for(lo.shape(), B.rw(), lo.write())->*[] __device__(size_t i, auto b, auto out) { + check_rw_view_types(); + out(i) = ::cuda::std::get<0>(b)(i) * ::cuda::std::get<1>(b)(i); + ::cuda::std::get<0>(b)(i) += 1.0; + }; + + // task with a read bundle dep: all views const + ctx.task(B.read(), lo.rw())->*[](cudaStream_t, auto b, auto) { + static_assert(::cuda::std::is_same_v<::cuda::std::tuple_element_t<0, decltype(b)>, slice>, + "read bundle view must const-qualify mutable fields"); + static_assert(::cuda::std::is_same_v<::cuda::std::tuple_element_t<1, decltype(b)>, slice>, + "read bundle view must const-qualify constant fields"); + }; + + // Mixed use: bundle dep and a bare dep on one of its own fields in the same task + ctx.parallel_for(lo.shape(), B.read(), lv.rw())->*[] __device__(size_t i, auto b, auto v2) { + v2(i) = ::cuda::std::get<0>(b)(i); + }; + + // host_launch with structured bindings on the bundle view; verify the results + ctx.host_launch(B.read(), lo.read())->*[N](auto b, auto out) { + auto& [v, ix] = b; + for (size_t i = 0; i < N; i++) + { + EXPECT(out(i) == 2.0 * ix(i)); + EXPECT(v(i) == 3.0); + } + }; + + ctx.finalize(); +} + +int main() +{ + // Creating constructor: fresh logical data from shapes, fields exposed + { + context ctx; + bundle>, field>> C(ctx, shape_of>(16), shape_of>(16)); + ctx.parallel_for(C.get_field<0>().shape(), C.write())->*[] __device__(size_t i, auto c) { + ::cuda::std::get<0>(c)(i) = 1.0; + ::cuda::std::get<1>(c)(i) = 2; + }; + // Bare-leaf consumption of a bundle-created field + ctx.host_launch(C.get_field<1>().read())->*[](auto ci) { + EXPECT(ci(0) == 2); + }; + ctx.finalize(); + } + + { + context ctx; + run(ctx); + } + + { + context ctx = graph_ctx(); + run(ctx); + } + + return 0; +} From f69e207cd193cb19f59121ca1f59ed72562579d5 Mon Sep 17 00:00:00 2001 From: Cedric AUGONNET Date: Sat, 15 Aug 2026 12:35:35 +0000 Subject: [PATCH 02/15] [STF] Bundles: token-aware grouping, launch and partitioner parallel_for overloads MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Tokens (void_interface deps) are filtered out of the lambda argument list by the constructs, so the regrouping adapter now computes USER-VISIBLE arities per slot (0 for token deps, per-non-void-field counts for bundles) and assembles the call with per-slot tuple pieces flattened by tuple_cat — zero-arity slots vanish, arity-1 slots stay forwarded references (reduction accumulators), bundles stay one tuple. Mixing tokens and bundles in one construct now works and is tested. Also adds the bundle-aware overloads for launch (spec+place, place, bare) and partitioner parallel_for on both context and backend_ctx. The ths-without-place launch form is omitted (ambiguous against the spec+place form; spell the place explicitly). Test additions: token+bundle parallel_for, launch with a bundle. Verified GB300 sm_103 / CUDA 13.4, c++17+c++20, stream+graph. Co-Authored-By: Claude Fable 5 --- .../__stf/internal/backend_ctx.cuh | 40 +++++++++ .../experimental/__stf/internal/bundle.cuh | 85 +++++++++++++------ .../experimental/__stf/internal/context.cuh | 40 +++++++++ cudax/test/stf/interface/bundle.cu | 15 ++++ 4 files changed, 154 insertions(+), 26 deletions(-) diff --git a/cudax/include/cuda/experimental/__stf/internal/backend_ctx.cuh b/cudax/include/cuda/experimental/__stf/internal/backend_ctx.cuh index 0e501e9c2304..3ac270ec38af 100644 --- a/cudax/include/cuda/experimental/__stf/internal/backend_ctx.cuh +++ b/cudax/include/cuda/experimental/__stf/internal/backend_ctx.cuh @@ -1235,6 +1235,46 @@ public: return parallel_for(self().default_exec_place(), mv(shape), mv(args)...); } + template < + typename partitioner_t, + typename exec_place_t, + typename S, + typename... Args, + ::cuda::std::enable_if_t<::cuda::std::is_base_of_v && reserved::any_bundle_dep_v, + int> = 0> + auto parallel_for(partitioner_t p, exec_place_t e_place, S shape, Args... args) + { + return reserved::make_bundle_scope( + [&](auto... flat) { + return this->parallel_for(mv(p), mv(e_place), mv(shape), mv(flat)...); + }, + mv(args)...); + } + + template >> + auto launch(thread_hierarchy_spec_t spec, exec_place e_place, Args... args) + { + return reserved::make_bundle_scope( + [&](auto... flat) { + return this->launch(mv(spec), mv(e_place), mv(flat)...); + }, + mv(args)...); + } + + template >> + auto launch(exec_place e_place, Args... args) + { + return launch(par(par()), mv(e_place), mv(args)...); + } + + template >> + auto launch(Args... args) + { + return launch(self().default_exec_place(), mv(args)...); + } + template >> auto host_launch(Args... args) { diff --git a/cudax/include/cuda/experimental/__stf/internal/bundle.cuh b/cudax/include/cuda/experimental/__stf/internal/bundle.cuh index f2b8cf1d6bc4..3e7fdfab36fb 100644 --- a/cudax/include/cuda/experimental/__stf/internal/bundle.cuh +++ b/cudax/include/cuda/experimental/__stf/internal/bundle.cuh @@ -62,6 +62,11 @@ class logical_data; template class shape_of; +class void_interface; + +template +class task_dep; + /** * @brief Trait tag marking a bundle field as constant (read-only ceiling). * @@ -132,13 +137,26 @@ inline constexpr bool is_bundle_dep_v = is_bundle_dep<::cuda::std::remove_cvref_ template inline constexpr bool any_bundle_dep_v = (is_bundle_dep_v || ... || false); -//! Number of flat dependencies one submitted argument expands to +//! Number of user-visible lambda arguments one flat dependency produces: +//! void_interface (token) dependencies are filtered out of the argument list +//! by the constructs and thus produce none. +template +struct dep_visible_arity : ::cuda::std::integral_constant +{}; + +template +struct dep_visible_arity> : ::cuda::std::integral_constant +{}; + +//! Number of user-visible lambda arguments one submitted argument produces template -struct slot_arity : ::cuda::std::integral_constant +struct visible_slot_arity + : ::cuda::std::integral_constant>::value> {}; template -struct slot_arity> : ::cuda::std::integral_constant +struct visible_slot_arity> + : ::cuda::std::integral_constant>::value + ... + 0)> {}; //! Expand one submitted dependency argument into a tuple of flat dependencies @@ -320,44 +338,56 @@ struct bundle_arg_adapter> return o; } - // Group the slot starting at flat position Offset (relative to the first - // flat view) out of the full forwarded argument tuple. NLead is the number - // of leading (non-dependency) arguments. + _CCCL_EXEC_CHECK_DISABLE + template + _CCCL_HOST_DEVICE static auto make_view_tuple(Tup& tup, ::std::index_sequence) + { + return ::cuda::std::tuple< + ::cuda::std::remove_cvref_t<::cuda::std::tuple_element_t>>...>( + ::cuda::std::get(tup)...); + } + + // One argument-tuple piece per slot, later flattened with tuple_cat: empty + // for zero-arity slots (token dependencies produce no user argument), a + // single forwarded reference for arity-1 slots (a reduction accumulator + // must remain a reference), a single by-value tuple of views for bundles. _CCCL_EXEC_CHECK_DISABLE template - _CCCL_HOST_DEVICE static decltype(auto) group_slot(Tup& tup) + _CCCL_HOST_DEVICE static auto slot_piece(Tup& tup) { constexpr size_t off = NLead + offset_of(Slot); - if constexpr (arity_of(Slot) == 1) + if constexpr (arity_of(Slot) == 0) + { + return ::cuda::std::tuple<>(); + } + else if constexpr (arity_of(Slot) == 1) { - // Forward the argument as-is (a reduction accumulator must remain a reference) - return static_cast<::cuda::std::tuple_element_t>&&>( - ::cuda::std::get(tup)); + return ::cuda::std::forward_as_tuple( + static_cast<::cuda::std::tuple_element_t>&&>(::cuda::std::get(tup))); } else { - return make_view_tuple(tup, ::std::make_index_sequence()); + return ::cuda::std::make_tuple(make_view_tuple(tup, ::std::make_index_sequence())); } } - _CCCL_EXEC_CHECK_DISABLE - template - _CCCL_HOST_DEVICE static auto make_view_tuple(Tup& tup, ::std::index_sequence) + // Type-level mirror of slot_piece, used to constrain operator() + template + using slot_piece_t = decltype(slot_piece(::cuda::std::declval())); + + template + static constexpr bool apply_invocable(::std::index_sequence) { - return ::cuda::std::tuple< - ::cuda::std::remove_cvref_t<::cuda::std::tuple_element_t>>...>( - ::cuda::std::get(tup)...); + return ::cuda::std::is_invocable_v...>; } - // Type-level mirror of group_slot, used to constrain operator() - template - using group_slot_t = decltype(group_slot(::cuda::std::declval())); - template static constexpr bool invocable_impl(::std::index_sequence, ::std::index_sequence) { - return ::cuda::std:: - is_invocable_v&&..., group_slot_t...>; + using args_tuple = decltype(::cuda::std::tuple_cat( + ::cuda::std::declval<::cuda::std::tuple<::cuda::std::tuple_element_t...>>(), + ::cuda::std::declval>()...)); + return apply_invocable(::std::make_index_sequence<::cuda::std::tuple_size_v>()); } template @@ -380,7 +410,10 @@ struct bundle_arg_adapter> _CCCL_HOST_DEVICE decltype(auto) call(Tup&& tup, ::std::index_sequence, ::std::index_sequence) { constexpr size_t n_lead = sizeof...(LeadIs); - return f(::cuda::std::get(mv(tup))..., group_slot(tup)...); + return ::cuda::std::apply( + f, + ::cuda::std::tuple_cat(::cuda::std::forward_as_tuple(::cuda::std::get(mv(tup))...), + slot_piece(tup)...)); } _CCCL_EXEC_CHECK_DISABLE @@ -459,7 +492,7 @@ private: template auto make_bundle_scope(MakeInner&& make_inner, Args... args) { - using arities = ::std::index_sequence::value...>; + using arities = ::std::index_sequence>::value...>; auto flat = ::std::tuple_cat(as_dep_tuple(mv(args))...); auto inner = ::std::apply(::cuda::std::forward(make_inner), mv(flat)); return bundle_scope(mv(inner)); diff --git a/cudax/include/cuda/experimental/__stf/internal/context.cuh b/cudax/include/cuda/experimental/__stf/internal/context.cuh index a2b0e0b1045d..7eb055e7238d 100644 --- a/cudax/include/cuda/experimental/__stf/internal/context.cuh +++ b/cudax/include/cuda/experimental/__stf/internal/context.cuh @@ -853,6 +853,46 @@ public: } #endif // !defined(CUDASTF_DISABLE_CODE_GENERATION) && _CCCL_CUDA_COMPILATION() + template < + typename partitioner_t, + typename exec_place_t, + typename S, + typename... Args, + ::cuda::std::enable_if_t<::cuda::std::is_base_of_v && reserved::any_bundle_dep_v, + int> = 0> + auto parallel_for(partitioner_t p, exec_place_t e_place, S shape, Args... args) + { + return reserved::make_bundle_scope( + [&](auto... flat) { + return this->parallel_for(mv(p), mv(e_place), mv(shape), mv(flat)...); + }, + mv(args)...); + } + + template >> + auto launch(thread_hierarchy_spec_t spec, exec_place e_place, Args... args) + { + return reserved::make_bundle_scope( + [&](auto... flat) { + return this->launch(mv(spec), mv(e_place), mv(flat)...); + }, + mv(args)...); + } + + template >> + auto launch(exec_place e_place, Args... args) + { + return launch(par(par()), mv(e_place), mv(args)...); + } + + template >> + auto launch(Args... args) + { + return launch(default_exec_place(), mv(args)...); + } + template >> auto host_launch(Args... args) { diff --git a/cudax/test/stf/interface/bundle.cu b/cudax/test/stf/interface/bundle.cu index 99e801690fd5..5f0c73b56036 100644 --- a/cudax/test/stf/interface/bundle.cu +++ b/cudax/test/stf/interface/bundle.cu @@ -78,6 +78,21 @@ void run(context& ctx) v2(i) = ::cuda::std::get<0>(b)(i); }; + // token dependency mixed with a bundle dependency: the token produces no + // lambda argument, and the bundle grouping must stay aligned + auto tok = ctx.token(); + ctx.parallel_for(lo.shape(), tok.rw(), B.read(), lo.rw())->*[] __device__(size_t i, auto b, auto out) { + out(i) += 0.0 * ::cuda::std::get<0>(b)(i); + }; + + // launch with a bundle dependency + ctx.launch(B.read(), lo.rw())->*[] __device__(auto th, auto b, auto out) { + for (size_t i = th.rank(); i < out.size(); i += th.size()) + { + out(i) += 0.0 * ::cuda::std::get<1>(b)(i); + } + }; + // host_launch with structured bindings on the bundle view; verify the results ctx.host_launch(B.read(), lo.read())->*[N](auto b, auto out) { auto& [v, ix] = b; From 06012e27f6454277d1c863ff426bfbab1e86333a Mon Sep 17 00:00:00 2001 From: Cedric AUGONNET Date: Sat, 15 Aug 2026 13:16:57 +0000 Subject: [PATCH 03/15] [STF] Bundles: guard bundle launch/partitioner overloads, cover cuda_kernel(_chain) CodeRabbit round: (1) the bundle partitioner parallel_for and launch overloads on context now live inside the same CUDASTF_DISABLE_CODE_GENERATION / CUDA-compilation guard as their non-bundle counterparts (verified by compiling with the macro defined); backend_ctx has no such guard convention, unchanged. (2) interface/bundle.cu now exercises cuda_kernel and cuda_kernel_chain with bundle dependencies on both backends. Co-Authored-By: Claude Fable 5 --- .../experimental/__stf/internal/context.cuh | 2 +- cudax/test/stf/interface/bundle.cu | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/cudax/include/cuda/experimental/__stf/internal/context.cuh b/cudax/include/cuda/experimental/__stf/internal/context.cuh index 7eb055e7238d..b564c2fd9fa8 100644 --- a/cudax/include/cuda/experimental/__stf/internal/context.cuh +++ b/cudax/include/cuda/experimental/__stf/internal/context.cuh @@ -851,7 +851,6 @@ public: { return parallel_for(default_exec_place(), mv(shape), mv(args)...); } -#endif // !defined(CUDASTF_DISABLE_CODE_GENERATION) && _CCCL_CUDA_COMPILATION() template < typename partitioner_t, @@ -892,6 +891,7 @@ public: { return launch(default_exec_place(), mv(args)...); } +#endif // !defined(CUDASTF_DISABLE_CODE_GENERATION) && _CCCL_CUDA_COMPILATION() template >> auto host_launch(Args... args) diff --git a/cudax/test/stf/interface/bundle.cu b/cudax/test/stf/interface/bundle.cu index 5f0c73b56036..98c2d09d4c39 100644 --- a/cudax/test/stf/interface/bundle.cu +++ b/cudax/test/stf/interface/bundle.cu @@ -26,6 +26,14 @@ using namespace cuda::experimental::stf; using test_bundle = bundle>, field, constant>>; +__global__ void bundle_kernel(slice v, slice ix, slice out) +{ + for (size_t i = threadIdx.x + blockIdx.x * blockDim.x; i < out.size(); i += blockDim.x * gridDim.x) + { + out(i) += 0.0 * (v(i) + ix(i)); + } +} + // The rw view must keep mutable fields mutable and const-qualify constant fields template __host__ __device__ void check_rw_view_types() @@ -93,6 +101,17 @@ void run(context& ctx) } }; + // cuda_kernel and cuda_kernel_chain with a bundle dependency + ctx.cuda_kernel(B.read(), lo.rw())->*[](auto b, auto out) { + return cuda_kernel_desc{bundle_kernel, 8, 32, 0, ::cuda::std::get<0>(b), ::cuda::std::get<1>(b), out}; + }; + + ctx.cuda_kernel_chain(B.read(), lo.rw())->*[](auto b, auto out) { + return ::std::vector{ + {bundle_kernel, 8, 32, 0, ::cuda::std::get<0>(b), ::cuda::std::get<1>(b), out}, + {bundle_kernel, 8, 32, 0, ::cuda::std::get<0>(b), ::cuda::std::get<1>(b), out}}; + }; + // host_launch with structured bindings on the bundle view; verify the results ctx.host_launch(B.read(), lo.read())->*[N](auto b, auto out) { auto& [v, ix] = b; From f534b2f2fbe49e6fbed33630098411670a8447ab Mon Sep 17 00:00:00 2001 From: Cedric AUGONNET Date: Sat, 15 Aug 2026 13:35:58 +0000 Subject: [PATCH 04/15] [STF] Bundles test: drop unused lambda capture (Clang14 -Werror) N is constexpr and read without odr-use; Clang14's -Wunused-lambda-capture rejects the explicit capture under -Werror (CTK12.0 Clang14 CI lane). Co-Authored-By: Claude Fable 5 --- cudax/test/stf/interface/bundle.cu | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cudax/test/stf/interface/bundle.cu b/cudax/test/stf/interface/bundle.cu index 98c2d09d4c39..f750c7a3f2ec 100644 --- a/cudax/test/stf/interface/bundle.cu +++ b/cudax/test/stf/interface/bundle.cu @@ -46,7 +46,7 @@ __host__ __device__ void check_rw_view_types() void run(context& ctx) { - const size_t N = 64; + constexpr size_t N = 64; ::std::vector vals(N, 2.0); ::std::vector idx(N); @@ -113,7 +113,7 @@ void run(context& ctx) }; // host_launch with structured bindings on the bundle view; verify the results - ctx.host_launch(B.read(), lo.read())->*[N](auto b, auto out) { + ctx.host_launch(B.read(), lo.read())->*[](auto b, auto out) { auto& [v, ix] = b; for (size_t i = 0; i < N; i++) { From 8d4e1fedb034aeba5c71574d8322699b9875b07e Mon Sep 17 00:00:00 2001 From: Cedric AUGONNET Date: Sat, 15 Aug 2026 13:55:11 +0000 Subject: [PATCH 05/15] [STF] Bundles: avoid identifier 'I' in headers (complex.h macro lint) The header-hygiene lane forbids bare 'I' in CCCL headers; bundle.cuh's field-index template parameters are now 'Idx'. Co-Authored-By: Claude Fable 5 --- .../experimental/__stf/internal/bundle.cuh | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/cudax/include/cuda/experimental/__stf/internal/bundle.cuh b/cudax/include/cuda/experimental/__stf/internal/bundle.cuh index 3e7fdfab36fb..a33c35814b1b 100644 --- a/cudax/include/cuda/experimental/__stf/internal/bundle.cuh +++ b/cudax/include/cuda/experimental/__stf/internal/bundle.cuh @@ -179,7 +179,7 @@ auto as_dep_tuple(T t) * * Bundles introduce no new ownership domain: they hold ordinary (refcounted) * logical data handles. Every field remains a first-class logical data - * retrievable with `get_field()`, so bundle-level and per-field task + * retrievable with `get_field()`, so bundle-level and per-field task * dependencies interoperate (they meet at the same logical data). * * @tparam Fields `field` descriptors, in canonical order @@ -192,8 +192,8 @@ class bundle public: static constexpr size_t n_fields = sizeof...(Fields); - template - using field_at = ::std::tuple_element_t>; + template + using field_at = ::std::tuple_element_t>; /** @brief Adopt existing logical data as the bundle's fields (no context needed) */ explicit bundle(logical_data... h) @@ -206,17 +206,17 @@ public: : handles(ctx.logical_data(mv(shapes))...) {} - /** @brief Access field `I` as its plain, first-class logical data */ - template + /** @brief Access field `Idx` as its plain, first-class logical data */ + template auto& get_field() { - return ::std::get(handles); + return ::std::get(handles); } - template + template const auto& get_field() const { - return ::std::get(handles); + return ::std::get(handles); } /** @brief Depend on every field with read access */ @@ -256,11 +256,11 @@ public: } private: - template + template auto leaf_dep() { - auto& ld = ::std::get(handles); - if constexpr (M == access_mode::read || field_at::is_constant) + auto& ld = ::std::get(handles); + if constexpr (M == access_mode::read || field_at::is_constant) { return ld.read(); } From ce5408eebc07d6ff18fbdf1245131b03c252da02 Mon Sep 17 00:00:00 2001 From: Cedric AUGONNET Date: Sat, 15 Aug 2026 15:54:39 +0000 Subject: [PATCH 06/15] [STF] Python bundles: same semantics as the C++ front end, zero shared code cuda.stf._experimental.bundles adds bundle / constant / bundle_dep / bundle_task over the flat bindings: bundle deps flatten into ordinary deps before the Cython task constructor, and bundle_task.get(i) regroups per-field views with each bundle counting as one slot (mirroring the numba_task interop wrapper idiom). The conformance semantics match the C++ side case for case: whole-bundle modes distribute per field (rw()/write() clamp constant fields to read), explicit requests above a field's ceiling raise, unspecified fields in per-field spellings default to read. Registering a CUDA-Array-Interface object now infers the device data place from cudaPointerGetAttributes (previously an opaque host-pinning assertion). tests/stf/test_bundles.py implements the shared conformance checklist: mode distribution, ceiling errors, slot counting, adoption vs registration, mixed bundle + bare-field use, device-array inference. Co-Authored-By: Claude Fable 5 --- .../cuda/stf/_experimental/__init__.py | 4 + .../cuda/stf/_experimental/bundles.py | 199 ++++++++++++++++++ python/cuda_stf/tests/stf/test_bundles.py | 99 +++++++++ 3 files changed, 302 insertions(+) create mode 100644 python/cuda_stf/cuda/stf/_experimental/bundles.py create mode 100644 python/cuda_stf/tests/stf/test_bundles.py diff --git a/python/cuda_stf/cuda/stf/_experimental/__init__.py b/python/cuda_stf/cuda/stf/_experimental/__init__.py index 49462fb664f1..be44194f1d82 100644 --- a/python/cuda_stf/cuda/stf/_experimental/__init__.py +++ b/python/cuda_stf/cuda/stf/_experimental/__init__.py @@ -31,6 +31,10 @@ _LAZY_SYMBOLS = { "AccessMode": "._stf_bindings", "CudaStream": "._stf_bindings", + "bundle": ".bundles", + "bundle_dep": ".bundles", + "bundle_task": ".bundles", + "constant": ".bundles", "async_resources": "._stf_bindings", "cond": "._stf_bindings", "context": "._stf_bindings", diff --git a/python/cuda_stf/cuda/stf/_experimental/bundles.py b/python/cuda_stf/cuda/stf/_experimental/bundles.py new file mode 100644 index 000000000000..7ec29c1da232 --- /dev/null +++ b/python/cuda_stf/cuda/stf/_experimental/bundles.py @@ -0,0 +1,199 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +"""Bundles: non-owning groups of logical data used as a single dependency. + +A :class:`bundle` ties several logical data together behind one object (the +three arrays of a CSR matrix, a graph's topology) so tasks can depend on the +whole group with a single argument, while every constituent ("field") remains +an ordinary logical data usable on its own. A bundle owns no data and no +dependency-tracking state: a bundle dependency expands into one ordinary +dependency per field before reaching the task, and :meth:`bundle_task.get` +reassembles the per-field views with the bundle counting as ONE slot. + +This mirrors the C++ ``bundle`` / ``field`` / ``constant`` feature +(see ``cudax/include/cuda/experimental/__stf/internal/bundle.cuh``); the two +front ends share no code but implement the same semantics: whole-bundle modes +distribute per field as the strongest mode the field admits (``rw()`` on a +bundle clamps ``constant`` fields to read), explicitly requesting more than a +field's ceiling raises, unspecified fields in per-field spellings default to +read, and one submitted dependency is one ``get`` slot. + +``constant`` is a promise about *users of this bundle*, not global +immutability: other views or bare handles may legitimately write the field, +and the ordinary read dependencies the bundle generates are what serialize +against those writers. + +Example +------- +>>> from cuda.stf._experimental.bundles import bundle, bundle_task, constant +>>> A = bundle(ctx, vals=vals_array, colind=constant(colind_array), +... rowptr=constant(rowptr_array)) +>>> with bundle_task(ctx, A.rw(), ly.rw()) as t: +... a = t.get(0) # namespace: a.vals, a.colind, a.rowptr (CAI views) +... y = t.get(1) # ordinary dependency: one slot each +""" + +from types import SimpleNamespace + +from cuda.stf._experimental import _stf_bindings as _b + +__all__ = ["bundle", "bundle_dep", "bundle_task", "constant"] + +_READ = _b.AccessMode.READ.value +_RW = _b.AccessMode.RW.value +_WRITE = _b.AccessMode.WRITE.value + + +class constant: + """Marks a bundle field as read-only through this bundle (a view ceiling).""" + + __slots__ = ("value",) + + def __init__(self, value): + self.value = value + + +def _register(ctx, value): + """Register an array-like, inferring the data place for device memory. + + CUDA Array Interface objects live on a device; registering them without a + device data place trips an opaque host-pinning assertion, so infer the + device from the pointer attributes. + """ + if isinstance(value, _b.logical_data): + return value + cai = getattr(value, "__cuda_array_interface__", None) + if cai is not None: + from cuda.bindings import runtime as _rt + + err, attr = _rt.cudaPointerGetAttributes(cai["data"][0]) + if int(err) == 0 and attr.type == _rt.cudaMemoryType.cudaMemoryTypeDevice: + return ctx.logical_data(value, _b.data_place.device(attr.device)) + return ctx.logical_data(value) + + +class bundle_dep: + """A group of ordinary deps submitted as a single argument (one slot).""" + + __slots__ = ("deps", "names") + + def __init__(self, deps, names): + self.deps = list(deps) + self.names = list(names) + + +class bundle: + """Non-owning group of logical data with named fields and ceilings. + + Field values may be existing :class:`logical_data` (adopted: handles are + shared, nothing is copied) or array-likes (registered). Wrapping a value + in :class:`constant` gives the field a read-only ceiling. Fields remain + first-class logical data, reachable as attributes (``b.vals``). + """ + + def __init__(self, ctx, **fields): + if not fields: + raise ValueError("a bundle needs at least one field") + self._names = [] + self._lds = {} + self._ceiling_read = set() + for name, value in fields.items(): + if isinstance(value, constant): + self._ceiling_read.add(name) + value = value.value + self._lds[name] = _register(ctx, value) + self._names.append(name) + + def __getattr__(self, name): + try: + return self._lds[name] + except KeyError: + raise AttributeError(name) from None + + def __len__(self): + return len(self._names) + + def read(self, dplace=None): + """Depend on every field with read access.""" + return self._make(dict.fromkeys(self._names, _READ), dplace) + + def rw(self, dplace=None): + """Depend on the bundle read-write: constant fields clamp to read.""" + return self._make(dict.fromkeys(self._names, _RW), dplace) + + def write(self, dplace=None): + """Depend on the bundle write-mode: constant fields clamp to read. + + Mutable fields are written (previous content discarded, not fetched); + constant fields are still fetched, since a writer typically needs + them to interpret what it writes. + """ + return self._make(dict.fromkeys(self._names, _WRITE), dplace) + + def dep(self, dplace=None, **modes): + """Per-field access modes; unspecified fields default to read. + + Explicitly requesting more than a constant field's ceiling raises + ``ValueError`` (distribution clamps; explicit excess is an error). + """ + m = dict.fromkeys(self._names, _READ) + for name, mode in modes.items(): + if name not in self._lds: + raise KeyError(f"bundle has no field {name!r}") + mode = int(mode) + if name in self._ceiling_read and mode != _READ: + raise ValueError( + f"field {name!r} is constant in this bundle (read-only ceiling)" + ) + m[name] = mode + return self._make(m, dplace) + + def _make(self, modes, dplace): + deps = [] + for name in self._names: + mode = _READ if name in self._ceiling_read else modes[name] + deps.append(_b.dep(self._lds[name], mode, dplace)) + return bundle_dep(deps, self._names) + + +class bundle_task: + """Context manager wrapping ``ctx.task``: flattens bundle dependencies and + regroups per-slot access, with each bundle counting as one slot. + + Non-bundle arguments pass through unchanged; ``get(i)`` returns the plain + CUDA Array Interface view for them, and a :class:`types.SimpleNamespace` + of per-field views for bundle slots. Other task methods (``stream_ptr``, + ``get_arg_cai``, ...) are forwarded to the underlying task. + """ + + def __init__(self, ctx, *args, **kwargs): + self._slots = [] # (arity, names or None) + flat = [] + for a in args: + if isinstance(a, bundle_dep): + self._slots.append((len(a.deps), a.names)) + flat.extend(a.deps) + else: + self._slots.append((1, None)) + flat.append(a) + self._task = ctx.task(*flat, **kwargs) + + def __enter__(self): + self._task.__enter__() + return self + + def __exit__(self, *exc): + return self._task.__exit__(*exc) + + def __getattr__(self, name): + return getattr(self._task, name) + + def get(self, slot): + """Per-slot view access; a bundle is one slot.""" + arity, names = self._slots[slot] + base = sum(s[0] for s in self._slots[:slot]) + if names is None: + return self._task.get_arg_cai(base) + views = [self._task.get_arg_cai(base + k) for k in range(arity)] + return SimpleNamespace(**dict(zip(names, views))) diff --git a/python/cuda_stf/tests/stf/test_bundles.py b/python/cuda_stf/tests/stf/test_bundles.py new file mode 100644 index 000000000000..7c3390539cff --- /dev/null +++ b/python/cuda_stf/tests/stf/test_bundles.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +"""Bundles: grouped logical data used as a single dependency. + +Mirrors the C++ conformance checklist (cudax/test/stf/interface/bundle.cu): +mode distribution against ceilings, loud error on explicit excess, +unspecified-fields-default-read, one submitted dependency = one slot, fields +remain first-class, mixed bundle + bare-field use in one task. +""" + +import numpy as np +import pytest + +import cuda.stf._experimental as stf +from cuda.stf._experimental import _stf_bindings as b +from cuda.stf._experimental.bundles import bundle, bundle_dep, bundle_task, constant + +READ = b.AccessMode.READ.value +RW = b.AccessMode.RW.value +WRITE = b.AccessMode.WRITE.value + + +def _modes(bd: bundle_dep): + return {name: d.mode for name, d in zip(bd.names, bd.deps)} + + +def test_bundle_modes_and_ceilings(): + ctx = b.context() + vals = np.zeros(8, dtype=np.float64) + idx = np.arange(8, dtype=np.int32) + B = bundle(ctx, vals=vals, idx=constant(idx)) + + # fields stay first-class logical data + assert isinstance(B.vals, b.logical_data) + assert len(B) == 2 + + # read: everything read + assert _modes(B.read()) == {"vals": READ, "idx": READ} + # rw distributes as the strongest admitted mode: constant clamps to read + assert _modes(B.rw()) == {"vals": RW, "idx": READ} + # write clamps constant fields to read (still fetched), mutable to write + assert _modes(B.write()) == {"vals": WRITE, "idx": READ} + # per-field spelling: unspecified fields default to read + assert _modes(B.dep(vals=RW)) == {"vals": RW, "idx": READ} + + # explicit excess over a ceiling raises; unknown fields raise + with pytest.raises(ValueError): + B.dep(idx=RW) + with pytest.raises(KeyError): + B.dep(nope=READ) + + ctx.finalize() + + +def test_bundle_task_slots(): + ctx = b.context() + vals = np.full(8, 2.0, dtype=np.float64) + idx = np.arange(8, dtype=np.int32) + out = np.zeros(8, dtype=np.float64) + + B = bundle(ctx, vals=vals, idx=constant(idx)) + lo = ctx.logical_data(out) + + # one bundle (2 fields) + one plain dep: bundle counts as ONE slot + with bundle_task(ctx, B.read(), lo.rw()) as t: + g = t.get(0) + assert set(vars(g)) == {"vals", "idx"} + v = g.vals.__cuda_array_interface__ + o = t.get(1).__cuda_array_interface__ + assert v["shape"] == (8,) and o["shape"] == (8,) + assert v["data"][0] != o["data"][0] + + # mixed use: bundle dep + bare dep on one of its own fields + with bundle_task(ctx, B.read(), B.vals.rw()) as t: + assert t.get(1).__cuda_array_interface__["shape"] == (8,) + + ctx.finalize() + + +def test_bundle_adopts_and_registers(): + ctx = b.context() + lv = ctx.logical_data(np.zeros(4, dtype=np.float32)) + B = bundle(ctx, vals=lv, aux=np.ones(4, dtype=np.float32)) + assert B.vals is lv # adopted, not re-registered + with bundle_task(ctx, B.rw()) as t: + assert set(vars(t.get(0))) == {"vals", "aux"} + ctx.finalize() + + +def test_bundle_device_array_inference(): + da = stf.DeviceArray(8, np.dtype(np.float32), b.data_place.device(0)) + ctx = b.context() + # registering a device CUDA-Array-Interface object must infer the device + # data place (no host-pinning assertion) + B = bundle(ctx, vals=da) + with bundle_task(ctx, B.rw()) as t: + assert t.get(0).vals.__cuda_array_interface__["shape"] == (8,) + ctx.finalize() From 825d6811a03e76799f747fb428a7d988d1e772f1 Mon Sep 17 00:00:00 2001 From: Cedric AUGONNET Date: Sat, 15 Aug 2026 16:00:19 +0000 Subject: [PATCH 07/15] [STF] Python bundles: sparse CG example (the port cg.py couldn't be) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit cg.py ported cg_csr_stackable.cu but 'simplified to use a dense matrix' — spelling three dependencies per task for a CSR matrix was the friction. cg_csr_bundle.py is the honest sparse port: the CSR matrix is one bundle (values tracked, structure constant), every task takes it as a single argument, numba kernels launch on the task stream, and the solve verifies A @ x == b. Co-Authored-By: Claude Fable 5 --- .../tests/stf/examples/cg_csr_bundle.py | 185 ++++++++++++++++++ 1 file changed, 185 insertions(+) create mode 100644 python/cuda_stf/tests/stf/examples/cg_csr_bundle.py diff --git a/python/cuda_stf/tests/stf/examples/cg_csr_bundle.py b/python/cuda_stf/tests/stf/examples/cg_csr_bundle.py new file mode 100644 index 000000000000..9c00f6d27ed0 --- /dev/null +++ b/python/cuda_stf/tests/stf/examples/cg_csr_bundle.py @@ -0,0 +1,185 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +"""Sparse conjugate gradient with the CSR matrix as ONE bundle dependency. + +Python port of ``cudax/examples/stf/linear_algebra/cg_csr.cu``. The earlier +``cg.py`` port simplified the problem to a dense matrix because a CSR matrix +is three arrays (values / column indices / row offsets) and every task had to +spell all three dependencies. A :class:`bundle` makes the sparse port direct: +the matrix is one object, tasks depend on it with a single argument, and the +structure arrays carry a read-only ceiling (``constant``) so a whole-bundle +``rw()`` could never write them. + +Every task below receives the matrix as one slot: ``t.get(0)`` yields a +namespace with ``vals`` / ``colind`` / ``rowptr`` views. +""" + +import numpy as np +import pytest + +pytest.importorskip("cuda.stf._experimental._stf_bindings") +numba = pytest.importorskip("numba") +from numba import cuda # noqa: E402 + +from cuda.stf._experimental import _stf_bindings as stf # noqa: E402 +from cuda.stf._experimental.bundles import bundle, bundle_task, constant # noqa: E402 + + +def _nb(view): + """Adapt a task view (CUDA Array Interface) to a numba device array.""" + return cuda.from_cuda_array_interface( + view.__cuda_array_interface__, owner=None, sync=False + ) + + +def _stream(t): + """The task's CUDA stream as a numba stream (kernels must run on it).""" + return cuda.external_stream(t.stream_ptr()) + + +@cuda.jit +def _spmv_kernel(vals, colind, rowptr, x, y): + i = cuda.grid(1) + if i < y.shape[0]: + acc = 0.0 + for k in range(rowptr[i], rowptr[i + 1]): + acc += vals[k] * x[colind[k]] + y[i] = acc + + +@cuda.jit +def _zero1(out): + out[0] = 0.0 + + +@cuda.jit +def _dot_kernel(a, b, out): + i = cuda.grid(1) + if i < a.shape[0]: + cuda.atomic.add(out, 0, a[i] * b[i]) + + +@cuda.jit +def _axpy_kernel(alpha_num, alpha_den, x, y): + """y += (alpha_num / alpha_den) * x""" + i = cuda.grid(1) + if i < y.shape[0]: + y[i] += (alpha_num[0] / alpha_den[0]) * x[i] + + +@cuda.jit +def _xpay_kernel(beta_num, beta_den, x, y): + """y = x + (beta_num / beta_den) * y""" + i = cuda.grid(1) + if i < y.shape[0]: + y[i] = x[i] + (beta_num[0] / beta_den[0]) * y[i] + + +def spmv(ctx, A, lx, ly, n): + """y = A @ x — the CSR matrix is a single dependency.""" + nb = (n + 127) // 128 + with bundle_task(ctx, A.read(), lx.read(), ly.rw()) as t: + a = t.get(0) + _spmv_kernel[nb, 128, _stream(t)]( + _nb(a.vals), _nb(a.colind), _nb(a.rowptr), _nb(t.get(1)), _nb(t.get(2)) + ) + + +def dot(ctx, la, lb, lres, n): + nb = (n + 127) // 128 + with bundle_task(ctx, la.read(), lb.read(), lres.rw()) as t: + s = _stream(t) + out = _nb(t.get(2)) + _zero1[1, 1, s](out) + _dot_kernel[nb, 128, s](_nb(t.get(0)), _nb(t.get(1)), out) + + +def test_cg_csr_bundle(): + n = 256 + + # Tridiagonal SPD system (2 on the diagonal, -1 off-diagonal) + rowptr, colind, vals = [0], [], [] + for i in range(n): + for j, v in ((i - 1, -1.0), (i, 2.0), (i + 1, -1.0)): + if 0 <= j < n: + colind.append(j) + vals.append(v) + rowptr.append(len(colind)) + + rng = np.random.default_rng(7) + b = rng.standard_normal(n) + + ctx = stf.context() + + # The matrix is one bundle: values tracked, structure read-only ceilinged + A = bundle( + ctx, + vals=np.array(vals, dtype=np.float64), + colind=constant(np.array(colind, dtype=np.int32)), + rowptr=constant(np.array(rowptr, dtype=np.int32)), + ) + + x_host = np.zeros(n) + lb = ctx.logical_data(b) + lx = ctx.logical_data(x_host) + lr = ctx.logical_data(b.copy()) # r = b - A@0 = b + lp = ctx.logical_data(b.copy()) # p = r + lap = ctx.logical_data(np.zeros(n)) + lrsold = ctx.logical_data(np.zeros(1)) + lrsnew = ctx.logical_data(np.zeros(1)) + lpap = ctx.logical_data(np.zeros(1)) + + nb = (n + 127) // 128 + dot(ctx, lr, lr, lrsold, n) + + for _ in range(2 * n): + spmv(ctx, A, lp, lap, n) + dot(ctx, lp, lap, lpap, n) + + with bundle_task(ctx, lrsold.read(), lpap.read(), lp.read(), lx.rw()) as t: + _axpy_kernel[nb, 128, _stream(t)]( + _nb(t.get(0)), _nb(t.get(1)), _nb(t.get(2)), _nb(t.get(3)) + ) + with bundle_task(ctx, lrsold.read(), lpap.read(), lap.read(), lr.rw()) as t: + # r -= alpha * Ap + _xpay_neg[nb, 128, _stream(t)]( + _nb(t.get(0)), _nb(t.get(1)), _nb(t.get(2)), _nb(t.get(3)) + ) + dot(ctx, lr, lr, lrsnew, n) + with bundle_task(ctx, lrsnew.read(), lrsold.read(), lr.read(), lp.rw()) as t: + _xpay_kernel[nb, 128, _stream(t)]( + _nb(t.get(0)), _nb(t.get(1)), _nb(t.get(2)), _nb(t.get(3)) + ) + with bundle_task(ctx, lrsnew.read(), lrsold.rw()) as t: + _copy1[1, 1, _stream(t)](_nb(t.get(0)), _nb(t.get(1))) + + ctx.finalize() + + # x_host received the write-back at finalize; verify A @ x ~= b + ax = np.zeros(n) + rp = np.array(rowptr) + ci = np.array(colind) + vv = np.array(vals) + xh = np.asarray(x_host) + for i in range(n): + ax[i] = np.dot(vv[rp[i] : rp[i + 1]], xh[ci[rp[i] : rp[i + 1]]]) + np.testing.assert_allclose(ax, b, atol=1e-6) + + +@cuda.jit +def _xpay_neg(alpha_num, alpha_den, x, y): + """y -= (alpha_num / alpha_den) * x""" + i = cuda.grid(1) + if i < y.shape[0]: + y[i] -= (alpha_num[0] / alpha_den[0]) * x[i] + + +@cuda.jit +def _copy1(src, dst): + dst[0] = src[0] + + +if __name__ == "__main__": + test_cg_csr_bundle() + print("cg_csr_bundle OK") From 27e9971722c9e8d542efab1d89c457dc806bf1a5 Mon Sep 17 00:00:00 2001 From: Cedric AUGONNET Date: Sat, 15 Aug 2026 16:01:03 +0000 Subject: [PATCH 08/15] [STF] Python bundles example: drop unused logical data (ruff F841) Co-Authored-By: Claude Fable 5 --- python/cuda_stf/tests/stf/examples/cg_csr_bundle.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/cuda_stf/tests/stf/examples/cg_csr_bundle.py b/python/cuda_stf/tests/stf/examples/cg_csr_bundle.py index 9c00f6d27ed0..f4ca9ae733e1 100644 --- a/python/cuda_stf/tests/stf/examples/cg_csr_bundle.py +++ b/python/cuda_stf/tests/stf/examples/cg_csr_bundle.py @@ -121,7 +121,6 @@ def test_cg_csr_bundle(): ) x_host = np.zeros(n) - lb = ctx.logical_data(b) lx = ctx.logical_data(x_host) lr = ctx.logical_data(b.copy()) # r = b - A@0 = b lp = ctx.logical_data(b.copy()) # p = r From 99f0222182e4c1970e0f6fc73e9a25bbb1f93b0a Mon Sep 17 00:00:00 2001 From: Cedric AUGONNET Date: Sat, 15 Aug 2026 16:09:30 +0000 Subject: [PATCH 09/15] [STF] Python bundles: ctx.task() accepts bundle deps directly MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review: the separate bundle_task wrapper forked the API for no reason. context.task now flattens bundle dependencies itself (several flat deps, one slot) and the task object gains get(slot) — uniform for plain and bundle dependencies alike: a plain dep's get(i) returns its CUDA Array Interface view, a bundle's returns a namespace of per-field views. bundle_task is gone; the example reads 'with ctx.task(A.read(), lx.read(), ly.rw()) as t'. Co-Authored-By: Claude Fable 5 --- .../cuda/stf/_experimental/__init__.py | 1 - .../stf/_experimental/_stf_bindings_impl.pyx | 30 +++++++++++ .../cuda/stf/_experimental/bundles.py | 54 +++---------------- .../tests/stf/examples/cg_csr_bundle.py | 14 ++--- python/cuda_stf/tests/stf/test_bundles.py | 12 ++--- 5 files changed, 49 insertions(+), 62 deletions(-) diff --git a/python/cuda_stf/cuda/stf/_experimental/__init__.py b/python/cuda_stf/cuda/stf/_experimental/__init__.py index be44194f1d82..e40aaaf9ed07 100644 --- a/python/cuda_stf/cuda/stf/_experimental/__init__.py +++ b/python/cuda_stf/cuda/stf/_experimental/__init__.py @@ -33,7 +33,6 @@ "CudaStream": "._stf_bindings", "bundle": ".bundles", "bundle_dep": ".bundles", - "bundle_task": ".bundles", "constant": ".bundles", "async_resources": "._stf_bindings", "cond": "._stf_bindings", diff --git a/python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx b/python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx index 349c234ca4a9..3c255aa29764 100644 --- a/python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx +++ b/python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx @@ -2133,6 +2133,9 @@ cdef class task: cdef _AliveFlag _alive # Grid rank of the exec place set through set_exec_place (1 = scalar) cdef int _grid_rank + # One entry per submitted dependency slot: (arity, field names or None). + # A bundle dependency contributes several flat deps but ONE slot. + cdef list _slot_map def __cinit__(self, context ctx): self._t = stf_task_create(ctx._ctx) @@ -2144,6 +2147,7 @@ cdef class task: self._mapper_states = [] self._alive = ctx._alive self._grid_rank = 1 + self._slot_map = [] def __dealloc__(self): # See stackable_logical_data.__dealloc__ for why a None _alive must @@ -2269,6 +2273,26 @@ cdef class task: n *= e return [self.get_stream_at_index(i) for i in range(n)] + def get(self, slot): + """Per-slot view access: one submitted dependency is one slot. + + For a plain dependency this returns its CUDA Array Interface view + (like ``get_arg_cai``); for a bundle dependency it returns a + ``types.SimpleNamespace`` with one view per field. + """ + if slot < 0 or slot >= len(self._slot_map): + raise IndexError(f"task has {len(self._slot_map)} dependency slots") + base = 0 + for s in range(slot): + base += self._slot_map[s][0] + arity, names = self._slot_map[slot] + if names is None: + return self.get_arg_cai(base) + from types import SimpleNamespace + + views = [self.get_arg_cai(base + k) for k in range(arity)] + return SimpleNamespace(**dict(zip(names, views))) + def get_arg(self, index) -> int: if self._lds_args[index]._is_token: raise RuntimeError("cannot materialize a token argument") @@ -3102,6 +3126,12 @@ cdef class context: for d in args: if isinstance(d, dep): t.add_dep(d) + t._slot_map.append((1, None)) + elif getattr(d, "_stf_bundle_dep", False): + # a bundle dependency: several flat deps, one slot + for leaf in d.deps: + t.add_dep(leaf) + t._slot_map.append((len(d.deps), list(d.names))) elif isinstance(d, exec_place): if exec_place_set: raise ValueError("Only one exec_place can be given") diff --git a/python/cuda_stf/cuda/stf/_experimental/bundles.py b/python/cuda_stf/cuda/stf/_experimental/bundles.py index 7ec29c1da232..babd1a691125 100644 --- a/python/cuda_stf/cuda/stf/_experimental/bundles.py +++ b/python/cuda_stf/cuda/stf/_experimental/bundles.py @@ -8,7 +8,7 @@ whole group with a single argument, while every constituent ("field") remains an ordinary logical data usable on its own. A bundle owns no data and no dependency-tracking state: a bundle dependency expands into one ordinary -dependency per field before reaching the task, and :meth:`bundle_task.get` +dependency per field before reaching the task, and ``task.get`` reassembles the per-field views with the bundle counting as ONE slot. This mirrors the C++ ``bundle`` / ``field`` / ``constant`` feature @@ -26,19 +26,17 @@ Example ------- ->>> from cuda.stf._experimental.bundles import bundle, bundle_task, constant +>>> from cuda.stf._experimental.bundles import bundle, constant >>> A = bundle(ctx, vals=vals_array, colind=constant(colind_array), ... rowptr=constant(rowptr_array)) ->>> with bundle_task(ctx, A.rw(), ly.rw()) as t: +>>> with ctx.task(A.rw(), ly.rw()) as t: ... a = t.get(0) # namespace: a.vals, a.colind, a.rowptr (CAI views) ... y = t.get(1) # ordinary dependency: one slot each """ -from types import SimpleNamespace - from cuda.stf._experimental import _stf_bindings as _b -__all__ = ["bundle", "bundle_dep", "bundle_task", "constant"] +__all__ = ["bundle", "bundle_dep", "constant"] _READ = _b.AccessMode.READ.value _RW = _b.AccessMode.RW.value @@ -76,6 +74,8 @@ def _register(ctx, value): class bundle_dep: """A group of ordinary deps submitted as a single argument (one slot).""" + _stf_bundle_dep = True + __slots__ = ("deps", "names") def __init__(self, deps, names): @@ -155,45 +155,3 @@ def _make(self, modes, dplace): mode = _READ if name in self._ceiling_read else modes[name] deps.append(_b.dep(self._lds[name], mode, dplace)) return bundle_dep(deps, self._names) - - -class bundle_task: - """Context manager wrapping ``ctx.task``: flattens bundle dependencies and - regroups per-slot access, with each bundle counting as one slot. - - Non-bundle arguments pass through unchanged; ``get(i)`` returns the plain - CUDA Array Interface view for them, and a :class:`types.SimpleNamespace` - of per-field views for bundle slots. Other task methods (``stream_ptr``, - ``get_arg_cai``, ...) are forwarded to the underlying task. - """ - - def __init__(self, ctx, *args, **kwargs): - self._slots = [] # (arity, names or None) - flat = [] - for a in args: - if isinstance(a, bundle_dep): - self._slots.append((len(a.deps), a.names)) - flat.extend(a.deps) - else: - self._slots.append((1, None)) - flat.append(a) - self._task = ctx.task(*flat, **kwargs) - - def __enter__(self): - self._task.__enter__() - return self - - def __exit__(self, *exc): - return self._task.__exit__(*exc) - - def __getattr__(self, name): - return getattr(self._task, name) - - def get(self, slot): - """Per-slot view access; a bundle is one slot.""" - arity, names = self._slots[slot] - base = sum(s[0] for s in self._slots[:slot]) - if names is None: - return self._task.get_arg_cai(base) - views = [self._task.get_arg_cai(base + k) for k in range(arity)] - return SimpleNamespace(**dict(zip(names, views))) diff --git a/python/cuda_stf/tests/stf/examples/cg_csr_bundle.py b/python/cuda_stf/tests/stf/examples/cg_csr_bundle.py index f4ca9ae733e1..7fce7b9bae3a 100644 --- a/python/cuda_stf/tests/stf/examples/cg_csr_bundle.py +++ b/python/cuda_stf/tests/stf/examples/cg_csr_bundle.py @@ -23,7 +23,7 @@ from numba import cuda # noqa: E402 from cuda.stf._experimental import _stf_bindings as stf # noqa: E402 -from cuda.stf._experimental.bundles import bundle, bundle_task, constant # noqa: E402 +from cuda.stf._experimental.bundles import bundle, constant # noqa: E402 def _nb(view): @@ -79,7 +79,7 @@ def _xpay_kernel(beta_num, beta_den, x, y): def spmv(ctx, A, lx, ly, n): """y = A @ x — the CSR matrix is a single dependency.""" nb = (n + 127) // 128 - with bundle_task(ctx, A.read(), lx.read(), ly.rw()) as t: + with ctx.task(A.read(), lx.read(), ly.rw()) as t: a = t.get(0) _spmv_kernel[nb, 128, _stream(t)]( _nb(a.vals), _nb(a.colind), _nb(a.rowptr), _nb(t.get(1)), _nb(t.get(2)) @@ -88,7 +88,7 @@ def spmv(ctx, A, lx, ly, n): def dot(ctx, la, lb, lres, n): nb = (n + 127) // 128 - with bundle_task(ctx, la.read(), lb.read(), lres.rw()) as t: + with ctx.task(la.read(), lb.read(), lres.rw()) as t: s = _stream(t) out = _nb(t.get(2)) _zero1[1, 1, s](out) @@ -136,21 +136,21 @@ def test_cg_csr_bundle(): spmv(ctx, A, lp, lap, n) dot(ctx, lp, lap, lpap, n) - with bundle_task(ctx, lrsold.read(), lpap.read(), lp.read(), lx.rw()) as t: + with ctx.task(lrsold.read(), lpap.read(), lp.read(), lx.rw()) as t: _axpy_kernel[nb, 128, _stream(t)]( _nb(t.get(0)), _nb(t.get(1)), _nb(t.get(2)), _nb(t.get(3)) ) - with bundle_task(ctx, lrsold.read(), lpap.read(), lap.read(), lr.rw()) as t: + with ctx.task(lrsold.read(), lpap.read(), lap.read(), lr.rw()) as t: # r -= alpha * Ap _xpay_neg[nb, 128, _stream(t)]( _nb(t.get(0)), _nb(t.get(1)), _nb(t.get(2)), _nb(t.get(3)) ) dot(ctx, lr, lr, lrsnew, n) - with bundle_task(ctx, lrsnew.read(), lrsold.read(), lr.read(), lp.rw()) as t: + with ctx.task(lrsnew.read(), lrsold.read(), lr.read(), lp.rw()) as t: _xpay_kernel[nb, 128, _stream(t)]( _nb(t.get(0)), _nb(t.get(1)), _nb(t.get(2)), _nb(t.get(3)) ) - with bundle_task(ctx, lrsnew.read(), lrsold.rw()) as t: + with ctx.task(lrsnew.read(), lrsold.rw()) as t: _copy1[1, 1, _stream(t)](_nb(t.get(0)), _nb(t.get(1))) ctx.finalize() diff --git a/python/cuda_stf/tests/stf/test_bundles.py b/python/cuda_stf/tests/stf/test_bundles.py index 7c3390539cff..57a0db26339c 100644 --- a/python/cuda_stf/tests/stf/test_bundles.py +++ b/python/cuda_stf/tests/stf/test_bundles.py @@ -14,7 +14,7 @@ import cuda.stf._experimental as stf from cuda.stf._experimental import _stf_bindings as b -from cuda.stf._experimental.bundles import bundle, bundle_dep, bundle_task, constant +from cuda.stf._experimental.bundles import bundle, bundle_dep, constant READ = b.AccessMode.READ.value RW = b.AccessMode.RW.value @@ -53,7 +53,7 @@ def test_bundle_modes_and_ceilings(): ctx.finalize() -def test_bundle_task_slots(): +def test_task_slots(): ctx = b.context() vals = np.full(8, 2.0, dtype=np.float64) idx = np.arange(8, dtype=np.int32) @@ -63,7 +63,7 @@ def test_bundle_task_slots(): lo = ctx.logical_data(out) # one bundle (2 fields) + one plain dep: bundle counts as ONE slot - with bundle_task(ctx, B.read(), lo.rw()) as t: + with ctx.task(B.read(), lo.rw()) as t: g = t.get(0) assert set(vars(g)) == {"vals", "idx"} v = g.vals.__cuda_array_interface__ @@ -72,7 +72,7 @@ def test_bundle_task_slots(): assert v["data"][0] != o["data"][0] # mixed use: bundle dep + bare dep on one of its own fields - with bundle_task(ctx, B.read(), B.vals.rw()) as t: + with ctx.task(B.read(), B.vals.rw()) as t: assert t.get(1).__cuda_array_interface__["shape"] == (8,) ctx.finalize() @@ -83,7 +83,7 @@ def test_bundle_adopts_and_registers(): lv = ctx.logical_data(np.zeros(4, dtype=np.float32)) B = bundle(ctx, vals=lv, aux=np.ones(4, dtype=np.float32)) assert B.vals is lv # adopted, not re-registered - with bundle_task(ctx, B.rw()) as t: + with ctx.task(B.rw()) as t: assert set(vars(t.get(0))) == {"vals", "aux"} ctx.finalize() @@ -94,6 +94,6 @@ def test_bundle_device_array_inference(): # registering a device CUDA-Array-Interface object must infer the device # data place (no host-pinning assertion) B = bundle(ctx, vals=da) - with bundle_task(ctx, B.rw()) as t: + with ctx.task(B.rw()) as t: assert t.get(0).vals.__cuda_array_interface__["shape"] == (8,) ctx.finalize() From 6e52392875a6dcc70625cd5cfc9b81695e12e27c Mon Sep 17 00:00:00 2001 From: Cedric AUGONNET Date: Sat, 15 Aug 2026 16:20:08 +0000 Subject: [PATCH 10/15] [STF] Python bundles: bundle views are namedtuples, one kernel argument MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review: passing three separate arrays to the SpMV kernel lost the bundle abstraction at the device boundary. task.get(slot) now returns a cached namedtuple instead of a namespace — numba types namedtuples of arrays, so the whole bundle view crosses the kernel boundary as ONE argument (the device-side analog of the C++ tuple view): @cuda.jit def _spmv_kernel(a, x, y): ... a.vals[k] * x[a.colind[k]] ... Also: tests and the example now use the public package surface (stf.AccessMode / stf.context / stf.data_place) instead of reaching into _stf_bindings, and compare IntFlag members directly instead of unwrapping .value. Co-Authored-By: Claude Fable 5 --- .../stf/_experimental/_stf_bindings_impl.pyx | 21 ++++++++++++++--- .../tests/stf/examples/cg_csr_bundle.py | 21 ++++++++++++----- python/cuda_stf/tests/stf/test_bundles.py | 23 +++++++++---------- 3 files changed, 44 insertions(+), 21 deletions(-) diff --git a/python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx b/python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx index 3c255aa29764..46fba4a3fd05 100644 --- a/python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx +++ b/python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx @@ -1070,6 +1070,20 @@ cdef class logical_data: ctx._alive = self._alive return ctx +_bundle_view_types = {} + + +def _bundle_view_type(names): + """Cached namedtuple type for a bundle's field names (task.get views).""" + t = _bundle_view_types.get(names) + if t is None: + from collections import namedtuple + + t = namedtuple("bundle_view", names) + _bundle_view_types[names] = t + return t + + class dep: __slots__ = ("ld", "mode", "dplace") # ld may be either a logical_data or a stackable_logical_data; both classes @@ -2288,10 +2302,11 @@ cdef class task: arity, names = self._slot_map[slot] if names is None: return self.get_arg_cai(base) - from types import SimpleNamespace - + # A namedtuple (rather than a plain namespace) so the whole view can + # cross kernel boundaries as ONE argument: numba typing supports + # namedtuples of arrays, preserving the bundle abstraction on device. views = [self.get_arg_cai(base + k) for k in range(arity)] - return SimpleNamespace(**dict(zip(names, views))) + return _bundle_view_type(tuple(names))(*views) def get_arg(self, index) -> int: if self._lds_args[index]._is_token: diff --git a/python/cuda_stf/tests/stf/examples/cg_csr_bundle.py b/python/cuda_stf/tests/stf/examples/cg_csr_bundle.py index 7fce7b9bae3a..89334c517842 100644 --- a/python/cuda_stf/tests/stf/examples/cg_csr_bundle.py +++ b/python/cuda_stf/tests/stf/examples/cg_csr_bundle.py @@ -22,7 +22,7 @@ numba = pytest.importorskip("numba") from numba import cuda # noqa: E402 -from cuda.stf._experimental import _stf_bindings as stf # noqa: E402 +import cuda.stf._experimental as stf # noqa: E402 from cuda.stf._experimental.bundles import bundle, constant # noqa: E402 @@ -38,13 +38,23 @@ def _stream(t): return cuda.external_stream(t.stream_ptr()) +def _nb_views(bundle_view): + """Adapt a whole bundle view: a namedtuple of numba device arrays. + + Numba types namedtuples of arrays, so the bundle stays ONE kernel + argument — the device-side analog of the C++ tuple view. + """ + return type(bundle_view)(*map(_nb, bundle_view)) + + @cuda.jit -def _spmv_kernel(vals, colind, rowptr, x, y): +def _spmv_kernel(a, x, y): + """The CSR matrix arrives as ONE argument (a namedtuple of arrays).""" i = cuda.grid(1) if i < y.shape[0]: acc = 0.0 - for k in range(rowptr[i], rowptr[i + 1]): - acc += vals[k] * x[colind[k]] + for k in range(a.rowptr[i], a.rowptr[i + 1]): + acc += a.vals[k] * x[a.colind[k]] y[i] = acc @@ -80,9 +90,8 @@ def spmv(ctx, A, lx, ly, n): """y = A @ x — the CSR matrix is a single dependency.""" nb = (n + 127) // 128 with ctx.task(A.read(), lx.read(), ly.rw()) as t: - a = t.get(0) _spmv_kernel[nb, 128, _stream(t)]( - _nb(a.vals), _nb(a.colind), _nb(a.rowptr), _nb(t.get(1)), _nb(t.get(2)) + _nb_views(t.get(0)), _nb(t.get(1)), _nb(t.get(2)) ) diff --git a/python/cuda_stf/tests/stf/test_bundles.py b/python/cuda_stf/tests/stf/test_bundles.py index 57a0db26339c..69862baec48f 100644 --- a/python/cuda_stf/tests/stf/test_bundles.py +++ b/python/cuda_stf/tests/stf/test_bundles.py @@ -13,12 +13,11 @@ import pytest import cuda.stf._experimental as stf -from cuda.stf._experimental import _stf_bindings as b from cuda.stf._experimental.bundles import bundle, bundle_dep, constant -READ = b.AccessMode.READ.value -RW = b.AccessMode.RW.value -WRITE = b.AccessMode.WRITE.value +READ = stf.AccessMode.READ +RW = stf.AccessMode.RW +WRITE = stf.AccessMode.WRITE def _modes(bd: bundle_dep): @@ -26,13 +25,13 @@ def _modes(bd: bundle_dep): def test_bundle_modes_and_ceilings(): - ctx = b.context() + ctx = stf.context() vals = np.zeros(8, dtype=np.float64) idx = np.arange(8, dtype=np.int32) B = bundle(ctx, vals=vals, idx=constant(idx)) # fields stay first-class logical data - assert isinstance(B.vals, b.logical_data) + assert hasattr(B.vals, "read") and hasattr(B.vals, "rw") assert len(B) == 2 # read: everything read @@ -54,7 +53,7 @@ def test_bundle_modes_and_ceilings(): def test_task_slots(): - ctx = b.context() + ctx = stf.context() vals = np.full(8, 2.0, dtype=np.float64) idx = np.arange(8, dtype=np.int32) out = np.zeros(8, dtype=np.float64) @@ -65,7 +64,7 @@ def test_task_slots(): # one bundle (2 fields) + one plain dep: bundle counts as ONE slot with ctx.task(B.read(), lo.rw()) as t: g = t.get(0) - assert set(vars(g)) == {"vals", "idx"} + assert set(g._fields) == {"vals", "idx"} v = g.vals.__cuda_array_interface__ o = t.get(1).__cuda_array_interface__ assert v["shape"] == (8,) and o["shape"] == (8,) @@ -79,18 +78,18 @@ def test_task_slots(): def test_bundle_adopts_and_registers(): - ctx = b.context() + ctx = stf.context() lv = ctx.logical_data(np.zeros(4, dtype=np.float32)) B = bundle(ctx, vals=lv, aux=np.ones(4, dtype=np.float32)) assert B.vals is lv # adopted, not re-registered with ctx.task(B.rw()) as t: - assert set(vars(t.get(0))) == {"vals", "aux"} + assert set(t.get(0)._fields) == {"vals", "aux"} ctx.finalize() def test_bundle_device_array_inference(): - da = stf.DeviceArray(8, np.dtype(np.float32), b.data_place.device(0)) - ctx = b.context() + da = stf.DeviceArray(8, np.dtype(np.float32), stf.data_place.device(0)) + ctx = stf.context() # registering a device CUDA-Array-Interface object must infer the device # data place (no host-pinning assertion) B = bundle(ctx, vals=da) From fa1269133a3b76871c8866c62237ee4dc532eb37 Mon Sep 17 00:00:00 2001 From: Cedric AUGONNET Date: Sat, 15 Aug 2026 16:24:11 +0000 Subject: [PATCH 11/15] [STF] Python bundles: ctx.bundle() constructor Review: bundle(ctx, **fields) was a method spelled as a free constructor. context.bundle(**fields) now mirrors ctx.logical_data, so the whole surface anchors on the context: A = ctx.bundle(vals=vals, colind=constant(colind)) with ctx.task(A.rw(), ly.rw()) as t: a = t.get(0) Co-Authored-By: Claude Fable 5 --- .../stf/_experimental/_stf_bindings_impl.pyx | 18 ++++++++++++++++++ .../cuda_stf/cuda/stf/_experimental/bundles.py | 6 +++--- .../tests/stf/examples/cg_csr_bundle.py | 5 ++--- python/cuda_stf/tests/stf/test_bundles.py | 10 +++++----- 4 files changed, 28 insertions(+), 11 deletions(-) diff --git a/python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx b/python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx index 46fba4a3fd05..603a78e1e144 100644 --- a/python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx +++ b/python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx @@ -3124,6 +3124,24 @@ cdef class context: def token(self): return logical_data.token(self) + def bundle(self, **fields): + """ + Create a :class:`~cuda.stf._experimental.bundles.bundle`: a non-owning + group of logical data usable as a single task dependency. + + Field values may be existing logical data (adopted) or array-likes + (registered); wrap a value in ``constant`` for a read-only ceiling. + + Example + ------- + >>> A = ctx.bundle(vals=vals, colind=constant(colind)) + >>> with ctx.task(A.rw(), ly.rw()) as t: + ... a = t.get(0) + """ + from cuda.stf._experimental.bundles import bundle as _bundle + + return _bundle(self, **fields) + def task(self, *args, symbol=None): """ Create a `task` diff --git a/python/cuda_stf/cuda/stf/_experimental/bundles.py b/python/cuda_stf/cuda/stf/_experimental/bundles.py index babd1a691125..3b69581e49ea 100644 --- a/python/cuda_stf/cuda/stf/_experimental/bundles.py +++ b/python/cuda_stf/cuda/stf/_experimental/bundles.py @@ -26,9 +26,9 @@ Example ------- ->>> from cuda.stf._experimental.bundles import bundle, constant ->>> A = bundle(ctx, vals=vals_array, colind=constant(colind_array), -... rowptr=constant(rowptr_array)) +>>> from cuda.stf._experimental import constant +>>> A = ctx.bundle(vals=vals_array, colind=constant(colind_array), +... rowptr=constant(rowptr_array)) >>> with ctx.task(A.rw(), ly.rw()) as t: ... a = t.get(0) # namespace: a.vals, a.colind, a.rowptr (CAI views) ... y = t.get(1) # ordinary dependency: one slot each diff --git a/python/cuda_stf/tests/stf/examples/cg_csr_bundle.py b/python/cuda_stf/tests/stf/examples/cg_csr_bundle.py index 89334c517842..c38c99cbb5c7 100644 --- a/python/cuda_stf/tests/stf/examples/cg_csr_bundle.py +++ b/python/cuda_stf/tests/stf/examples/cg_csr_bundle.py @@ -23,7 +23,7 @@ from numba import cuda # noqa: E402 import cuda.stf._experimental as stf # noqa: E402 -from cuda.stf._experimental.bundles import bundle, constant # noqa: E402 +from cuda.stf._experimental.bundles import constant # noqa: E402 def _nb(view): @@ -122,8 +122,7 @@ def test_cg_csr_bundle(): ctx = stf.context() # The matrix is one bundle: values tracked, structure read-only ceilinged - A = bundle( - ctx, + A = ctx.bundle( vals=np.array(vals, dtype=np.float64), colind=constant(np.array(colind, dtype=np.int32)), rowptr=constant(np.array(rowptr, dtype=np.int32)), diff --git a/python/cuda_stf/tests/stf/test_bundles.py b/python/cuda_stf/tests/stf/test_bundles.py index 69862baec48f..d30fc8207021 100644 --- a/python/cuda_stf/tests/stf/test_bundles.py +++ b/python/cuda_stf/tests/stf/test_bundles.py @@ -13,7 +13,7 @@ import pytest import cuda.stf._experimental as stf -from cuda.stf._experimental.bundles import bundle, bundle_dep, constant +from cuda.stf._experimental.bundles import bundle_dep, constant READ = stf.AccessMode.READ RW = stf.AccessMode.RW @@ -28,7 +28,7 @@ def test_bundle_modes_and_ceilings(): ctx = stf.context() vals = np.zeros(8, dtype=np.float64) idx = np.arange(8, dtype=np.int32) - B = bundle(ctx, vals=vals, idx=constant(idx)) + B = ctx.bundle(vals=vals, idx=constant(idx)) # fields stay first-class logical data assert hasattr(B.vals, "read") and hasattr(B.vals, "rw") @@ -58,7 +58,7 @@ def test_task_slots(): idx = np.arange(8, dtype=np.int32) out = np.zeros(8, dtype=np.float64) - B = bundle(ctx, vals=vals, idx=constant(idx)) + B = ctx.bundle(vals=vals, idx=constant(idx)) lo = ctx.logical_data(out) # one bundle (2 fields) + one plain dep: bundle counts as ONE slot @@ -80,7 +80,7 @@ def test_task_slots(): def test_bundle_adopts_and_registers(): ctx = stf.context() lv = ctx.logical_data(np.zeros(4, dtype=np.float32)) - B = bundle(ctx, vals=lv, aux=np.ones(4, dtype=np.float32)) + B = ctx.bundle(vals=lv, aux=np.ones(4, dtype=np.float32)) assert B.vals is lv # adopted, not re-registered with ctx.task(B.rw()) as t: assert set(t.get(0)._fields) == {"vals", "aux"} @@ -92,7 +92,7 @@ def test_bundle_device_array_inference(): ctx = stf.context() # registering a device CUDA-Array-Interface object must infer the device # data place (no host-pinning assertion) - B = bundle(ctx, vals=da) + B = ctx.bundle(vals=da) with ctx.task(B.rw()) as t: assert t.get(0).vals.__cuda_array_interface__["shape"] == (8,) ctx.finalize() From 9df192ef345d3e8add0e54bbe9817f122b029550 Mon Sep 17 00:00:00 2001 From: Cedric AUGONNET Date: Sat, 15 Aug 2026 16:31:12 +0000 Subject: [PATCH 12/15] [STF] Python bundles: fields must be logical data MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review: passing raw arrays (vals=np.zeros(...)) conflated registration with grouping and hid the data-place policy inside the bundle (which is what forced the pointer-attribute inference). Bundles now group handles only — register arrays first with ctx.logical_data, which owns the placement policy explicitly. Raw values raise a TypeError pointing there. Duck-typed on read()/rw() so stackable logical data qualify. Co-Authored-By: Claude Fable 5 --- .../stf/_experimental/_stf_bindings_impl.pyx | 5 ++- .../cuda/stf/_experimental/bundles.py | 37 ++++++++----------- .../tests/stf/examples/cg_csr_bundle.py | 6 +-- python/cuda_stf/tests/stf/test_bundles.py | 27 ++++++++------ 4 files changed, 36 insertions(+), 39 deletions(-) diff --git a/python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx b/python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx index 603a78e1e144..cad21a99f4a6 100644 --- a/python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx +++ b/python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx @@ -3129,8 +3129,9 @@ cdef class context: Create a :class:`~cuda.stf._experimental.bundles.bundle`: a non-owning group of logical data usable as a single task dependency. - Field values may be existing logical data (adopted) or array-likes - (registered); wrap a value in ``constant`` for a read-only ceiling. + Field values must be logical data (bundles group handles; register + arrays first with ``ctx.logical_data``). Wrap a value in + ``constant`` for a read-only ceiling. Example ------- diff --git a/python/cuda_stf/cuda/stf/_experimental/bundles.py b/python/cuda_stf/cuda/stf/_experimental/bundles.py index 3b69581e49ea..f597cf7ab02a 100644 --- a/python/cuda_stf/cuda/stf/_experimental/bundles.py +++ b/python/cuda_stf/cuda/stf/_experimental/bundles.py @@ -52,23 +52,16 @@ def __init__(self, value): self.value = value -def _register(ctx, value): - """Register an array-like, inferring the data place for device memory. - - CUDA Array Interface objects live on a device; registering them without a - device data place trips an opaque host-pinning assertion, so infer the - device from the pointer attributes. - """ - if isinstance(value, _b.logical_data): - return value - cai = getattr(value, "__cuda_array_interface__", None) - if cai is not None: - from cuda.bindings import runtime as _rt - - err, attr = _rt.cudaPointerGetAttributes(cai["data"][0]) - if int(err) == 0 and attr.type == _rt.cudaMemoryType.cudaMemoryTypeDevice: - return ctx.logical_data(value, _b.data_place.device(attr.device)) - return ctx.logical_data(value) +def _check_field(name, value): + """Bundle fields must be logical data: bundles group handles, they do not + register data (that is ``ctx.logical_data``'s job, with its explicit data + place policy). Duck-typed so stackable logical data qualify too.""" + if not (hasattr(value, "read") and hasattr(value, "rw")): + raise TypeError( + f"bundle field {name!r} must be a logical data (register the array " + "first with ctx.logical_data(...))" + ) + return value class bundle_dep: @@ -86,10 +79,10 @@ def __init__(self, deps, names): class bundle: """Non-owning group of logical data with named fields and ceilings. - Field values may be existing :class:`logical_data` (adopted: handles are - shared, nothing is copied) or array-likes (registered). Wrapping a value - in :class:`constant` gives the field a read-only ceiling. Fields remain - first-class logical data, reachable as attributes (``b.vals``). + Field values must be logical data (handles are shared, nothing is copied + or registered — register arrays first with ``ctx.logical_data``). + Wrapping a value in :class:`constant` gives the field a read-only + ceiling. Fields remain first-class, reachable as attributes (``b.vals``). """ def __init__(self, ctx, **fields): @@ -102,7 +95,7 @@ def __init__(self, ctx, **fields): if isinstance(value, constant): self._ceiling_read.add(name) value = value.value - self._lds[name] = _register(ctx, value) + self._lds[name] = _check_field(name, value) self._names.append(name) def __getattr__(self, name): diff --git a/python/cuda_stf/tests/stf/examples/cg_csr_bundle.py b/python/cuda_stf/tests/stf/examples/cg_csr_bundle.py index c38c99cbb5c7..30eb7173be42 100644 --- a/python/cuda_stf/tests/stf/examples/cg_csr_bundle.py +++ b/python/cuda_stf/tests/stf/examples/cg_csr_bundle.py @@ -123,9 +123,9 @@ def test_cg_csr_bundle(): # The matrix is one bundle: values tracked, structure read-only ceilinged A = ctx.bundle( - vals=np.array(vals, dtype=np.float64), - colind=constant(np.array(colind, dtype=np.int32)), - rowptr=constant(np.array(rowptr, dtype=np.int32)), + vals=ctx.logical_data(np.array(vals, dtype=np.float64)), + colind=constant(ctx.logical_data(np.array(colind, dtype=np.int32))), + rowptr=constant(ctx.logical_data(np.array(rowptr, dtype=np.int32))), ) x_host = np.zeros(n) diff --git a/python/cuda_stf/tests/stf/test_bundles.py b/python/cuda_stf/tests/stf/test_bundles.py index d30fc8207021..44d9679f01a5 100644 --- a/python/cuda_stf/tests/stf/test_bundles.py +++ b/python/cuda_stf/tests/stf/test_bundles.py @@ -26,8 +26,8 @@ def _modes(bd: bundle_dep): def test_bundle_modes_and_ceilings(): ctx = stf.context() - vals = np.zeros(8, dtype=np.float64) - idx = np.arange(8, dtype=np.int32) + vals = ctx.logical_data(np.zeros(8, dtype=np.float64)) + idx = ctx.logical_data(np.arange(8, dtype=np.int32)) B = ctx.bundle(vals=vals, idx=constant(idx)) # fields stay first-class logical data @@ -54,8 +54,8 @@ def test_bundle_modes_and_ceilings(): def test_task_slots(): ctx = stf.context() - vals = np.full(8, 2.0, dtype=np.float64) - idx = np.arange(8, dtype=np.int32) + vals = ctx.logical_data(np.full(8, 2.0, dtype=np.float64)) + idx = ctx.logical_data(np.arange(8, dtype=np.int32)) out = np.zeros(8, dtype=np.float64) B = ctx.bundle(vals=vals, idx=constant(idx)) @@ -77,22 +77,25 @@ def test_task_slots(): ctx.finalize() -def test_bundle_adopts_and_registers(): +def test_bundle_adopts_handles_only(): ctx = stf.context() lv = ctx.logical_data(np.zeros(4, dtype=np.float32)) - B = ctx.bundle(vals=lv, aux=np.ones(4, dtype=np.float32)) - assert B.vals is lv # adopted, not re-registered + la = ctx.logical_data(np.ones(4, dtype=np.float32)) + B = ctx.bundle(vals=lv, aux=la) + assert B.vals is lv # adopted: the same handle, nothing copied with ctx.task(B.rw()) as t: assert set(t.get(0)._fields) == {"vals", "aux"} + # bundles group handles, they do not register arrays + with pytest.raises(TypeError): + ctx.bundle(vals=np.zeros(4)) ctx.finalize() -def test_bundle_device_array_inference(): - da = stf.DeviceArray(8, np.dtype(np.float32), stf.data_place.device(0)) +def test_bundle_device_logical_data(): + dplace = stf.data_place.device(0) + da = stf.DeviceArray(8, np.dtype(np.float32), dplace) ctx = stf.context() - # registering a device CUDA-Array-Interface object must infer the device - # data place (no host-pinning assertion) - B = ctx.bundle(vals=da) + B = ctx.bundle(vals=ctx.logical_data(da, dplace)) with ctx.task(B.rw()) as t: assert t.get(0).vals.__cuda_array_interface__["shape"] == (8,) ctx.finalize() From d30994ebecfe0358bdd025ce742ec9fc7218b0b9 Mon Sep 17 00:00:00 2001 From: Cedric AUGONNET Date: Sat, 15 Aug 2026 17:29:51 +0000 Subject: [PATCH 13/15] [STF] Python bundles: document when NOT to use a bundle MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bundles model objects (invariant-bound fields, library descriptors), not loose collections: a solver workspace whose tasks touch varying subsets with varying modes (Krylov vectors) should keep bare deps — grouping would over-declare and serialize tasks that share no data. Co-Authored-By: Claude Fable 5 --- python/cuda_stf/cuda/stf/_experimental/bundles.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/python/cuda_stf/cuda/stf/_experimental/bundles.py b/python/cuda_stf/cuda/stf/_experimental/bundles.py index f597cf7ab02a..0c60d4427f4f 100644 --- a/python/cuda_stf/cuda/stf/_experimental/bundles.py +++ b/python/cuda_stf/cuda/stf/_experimental/bundles.py @@ -19,6 +19,15 @@ field's ceiling raises, unspecified fields in per-field spellings default to read, and one submitted dependency is one ``get`` slot. +When to use a bundle — and when not to: bundles model data that is one +OBJECT at the level users reason about (a sparse matrix's invariant-bound +arrays, a graph's topology, a mesh's coordinates and connectivity, a library +descriptor's constituents). They are not a dependency-count reducer for +loosely related arrays: a solver workspace whose tasks touch different +subsets with different modes each time (e.g. Krylov vectors R/P/V/S/T) +should keep bare per-array dependencies — grouping such arrays would +over-declare and serialize tasks that share no data. + ``constant`` is a promise about *users of this bundle*, not global immutability: other views or bare handles may legitimately write the field, and the ordinary read dependencies the bundle generates are what serialize From a7c3c4e0b260b3748b7abd14e36bd7660455f9e5 Mon Sep 17 00:00:00 2001 From: Cedric AUGONNET Date: Sat, 15 Aug 2026 17:36:07 +0000 Subject: [PATCH 14/15] [STF] Python bundles: AccessMode.NONE excludes a field; docs correction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review: bundling adds no synchronization of its own — a bundle dep flattens to exactly the leaf deps you would write by hand. The docs now say so precisely: the cost of whole-object spellings and the read-default is acquiring (ordering + transferring) fields a task never touches — right for objects, meaningless for workspaces. AccessMode.NONE completes the mode set: an excluded field contributes no dependency and no transfer, and its view is None (the namedtuple keeps the field for shape stability). Also legitimate for object use: structure-only access to a matrix without acquiring its values. Co-Authored-By: Claude Fable 5 --- .../stf/_experimental/_stf_bindings_impl.pyx | 16 ++++++-- .../cuda/stf/_experimental/bundles.py | 37 +++++++++++++++---- python/cuda_stf/tests/stf/test_bundles.py | 10 ++++- 3 files changed, 50 insertions(+), 13 deletions(-) diff --git a/python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx b/python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx index cad21a99f4a6..b2fe68ffcdfe 100644 --- a/python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx +++ b/python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx @@ -2299,13 +2299,21 @@ cdef class task: base = 0 for s in range(slot): base += self._slot_map[s][0] - arity, names = self._slot_map[slot] + arity, names, acquired = self._slot_map[slot] if names is None: return self.get_arg_cai(base) # A namedtuple (rather than a plain namespace) so the whole view can # cross kernel boundaries as ONE argument: numba typing supports # namedtuples of arrays, preserving the bundle abstraction on device. - views = [self.get_arg_cai(base + k) for k in range(arity)] + # Fields excluded with mode NONE are present but None. + views = [] + k = 0 + for got in acquired: + if got: + views.append(self.get_arg_cai(base + k)) + k += 1 + else: + views.append(None) return _bundle_view_type(tuple(names))(*views) def get_arg(self, index) -> int: @@ -3160,12 +3168,12 @@ cdef class context: for d in args: if isinstance(d, dep): t.add_dep(d) - t._slot_map.append((1, None)) + t._slot_map.append((1, None, None)) elif getattr(d, "_stf_bundle_dep", False): # a bundle dependency: several flat deps, one slot for leaf in d.deps: t.add_dep(leaf) - t._slot_map.append((len(d.deps), list(d.names))) + t._slot_map.append((len(d.deps), list(d.names), list(d.acquired))) elif isinstance(d, exec_place): if exec_place_set: raise ValueError("Only one exec_place can be given") diff --git a/python/cuda_stf/cuda/stf/_experimental/bundles.py b/python/cuda_stf/cuda/stf/_experimental/bundles.py index 0c60d4427f4f..8be57ba4ff80 100644 --- a/python/cuda_stf/cuda/stf/_experimental/bundles.py +++ b/python/cuda_stf/cuda/stf/_experimental/bundles.py @@ -25,8 +25,13 @@ descriptor's constituents). They are not a dependency-count reducer for loosely related arrays: a solver workspace whose tasks touch different subsets with different modes each time (e.g. Krylov vectors R/P/V/S/T) -should keep bare per-array dependencies — grouping such arrays would -over-declare and serialize tasks that share no data. +should keep bare per-array dependencies. Bundling adds no synchronization +of its own (a bundle dependency flattens to exactly the leaf dependencies +you would write by hand) — but the whole-object spellings and the +read-default acquire every field, which for a workspace means transfers +and ordering against fields the task never touches, for no modeling gain. +Use ``AccessMode.NONE`` to exclude a field where object-level access is +otherwise right. ``constant`` is a promise about *users of this bundle*, not global immutability: other views or bare handles may legitimately write the field, @@ -47,6 +52,7 @@ __all__ = ["bundle", "bundle_dep", "constant"] +_NONE = _b.AccessMode.NONE.value _READ = _b.AccessMode.READ.value _RW = _b.AccessMode.RW.value _WRITE = _b.AccessMode.WRITE.value @@ -74,15 +80,21 @@ def _check_field(name, value): class bundle_dep: - """A group of ordinary deps submitted as a single argument (one slot).""" + """A group of ordinary deps submitted as a single argument (one slot). + + ``names`` covers every field of the bundle; ``acquired`` marks which of + them contribute a dependency (fields excluded with mode ``NONE`` are not + acquired at all — their view is ``None``). + """ _stf_bundle_dep = True - __slots__ = ("deps", "names") + __slots__ = ("deps", "names", "acquired") - def __init__(self, deps, names): + def __init__(self, deps, names, acquired): self.deps = list(deps) self.names = list(names) + self.acquired = list(acquired) class bundle: @@ -138,13 +150,15 @@ def dep(self, dplace=None, **modes): Explicitly requesting more than a constant field's ceiling raises ``ValueError`` (distribution clamps; explicit excess is an error). + A field set to ``AccessMode.NONE`` is not acquired at all: no + dependency, no transfer, and its view is ``None``. """ m = dict.fromkeys(self._names, _READ) for name, mode in modes.items(): if name not in self._lds: raise KeyError(f"bundle has no field {name!r}") mode = int(mode) - if name in self._ceiling_read and mode != _READ: + if name in self._ceiling_read and mode not in (_READ, _NONE): raise ValueError( f"field {name!r} is constant in this bundle (read-only ceiling)" ) @@ -153,7 +167,14 @@ def dep(self, dplace=None, **modes): def _make(self, modes, dplace): deps = [] + acquired = [] for name in self._names: - mode = _READ if name in self._ceiling_read else modes[name] + mode = modes[name] + if mode != _NONE and name in self._ceiling_read: + mode = _READ + if mode == _NONE: + acquired.append(False) + continue + acquired.append(True) deps.append(_b.dep(self._lds[name], mode, dplace)) - return bundle_dep(deps, self._names) + return bundle_dep(deps, self._names, acquired) diff --git a/python/cuda_stf/tests/stf/test_bundles.py b/python/cuda_stf/tests/stf/test_bundles.py index 44d9679f01a5..23a5f44cb47f 100644 --- a/python/cuda_stf/tests/stf/test_bundles.py +++ b/python/cuda_stf/tests/stf/test_bundles.py @@ -21,7 +21,8 @@ def _modes(bd: bundle_dep): - return {name: d.mode for name, d in zip(bd.names, bd.deps)} + acquired_names = [n for n, got in zip(bd.names, bd.acquired) if got] + return {name: d.mode for name, d in zip(acquired_names, bd.deps)} def test_bundle_modes_and_ceilings(): @@ -49,6 +50,13 @@ def test_bundle_modes_and_ceilings(): with pytest.raises(KeyError): B.dep(nope=READ) + # NONE excludes a field: no dep, no transfer, view is None + bd = B.dep(vals=RW, idx=stf.AccessMode.NONE) + assert _modes(bd) == {"vals": RW} + with ctx.task(bd) as t: + g = t.get(0) + assert g.idx is None and g.vals is not None + ctx.finalize() From d04ac17ec483c502f74f87ae2f72146151939c07 Mon Sep 17 00:00:00 2001 From: Cedric AUGONNET Date: Sun, 16 Aug 2026 08:56:37 +0000 Subject: [PATCH 15/15] [STF] Python bundles: stackable contexts get the same integration MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit stackable_context.bundle() / stackable_context.task() accepting bundle deps / stackable_task.get(slot) — the symmetric counterparts of the plain-context integration, so bundles compose with graph scopes and launchable_graph_scope record/replay. Exercised by a Newton physics phase-x-group prototype where per-State token bundles ({rigid, soft}) replace hand-maintained token dicts: whole-bundle modes distribute, sensors use dep(rigid=READ, soft=NONE), and the recorded task graph replays at parity with native capture while overlapping sensor work the coarse-token layout serializes. Co-Authored-By: Claude Fable 5 --- .../stf/_experimental/_stf_bindings_impl.pyx | 41 +++++++++++++++++++ python/cuda_stf/tests/stf/test_bundles.py | 21 ++++++++++ 2 files changed, 62 insertions(+) diff --git a/python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx b/python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx index b2fe68ffcdfe..039ab381d3f3 100644 --- a/python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx +++ b/python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx @@ -3504,6 +3504,9 @@ cdef class stackable_logical_data: cdef class stackable_task: cdef stf_task_handle _t cdef stf_ctx_handle _ctx + # One entry per submitted dependency slot: (arity, field names or None, + # acquired flags or None). A bundle dependency is ONE slot. + cdef list _slot_map cdef list _lds_args # Retain exec places and per-dep data-place overrides referenced by the task. cdef list _owners @@ -3514,6 +3517,7 @@ cdef class stackable_task: cdef _AliveFlag _alive def __cinit__(self, stackable_context ctx): + self._slot_map = [] self._t = stf_stackable_task_create(ctx._ctx) if self._t == NULL: raise RuntimeError("failed to create STF stackable task") @@ -3602,6 +3606,31 @@ cdef class stackable_task: cdef void *ptr = stf_task_get(self._t, index) return ptr + def get(self, slot): + """Per-slot view access: one submitted dependency is one slot. + + Mirrors :meth:`task.get`: plain dependencies return their CUDA Array + Interface view; bundle dependencies return a namedtuple of per-field + views (excluded fields are None). + """ + if slot < 0 or slot >= len(self._slot_map): + raise IndexError(f"task has {len(self._slot_map)} dependency slots") + base = 0 + for i in range(slot): + base += self._slot_map[i][0] + arity, names, acquired = self._slot_map[slot] + if names is None: + return self.get_arg_cai(base) + views = [] + k = 0 + for got in acquired: + if got: + views.append(self.get_arg_cai(base + k)) + k += 1 + else: + views.append(None) + return _bundle_view_type(tuple(names))(*views) + def get_arg_cai(self, index): """Return the argument as a CUDA Array Interface v3 object. @@ -4497,6 +4526,12 @@ cdef class stackable_context: raise RuntimeError("failed to create stackable token") return out + def bundle(self, **fields): + """Create a bundle over stackable logical data (see :meth:`context.bundle`).""" + from cuda.stf._experimental.bundles import bundle as _bundle + + return _bundle(self, **fields) + def task(self, *args, symbol=None): """Create a task on the head (innermost) scope of this context.""" exec_place_set = False @@ -4506,6 +4541,12 @@ cdef class stackable_context: for d in args: if isinstance(d, dep): t.add_dep(d) + t._slot_map.append((1, None, None)) + elif getattr(d, "_stf_bundle_dep", False): + # a bundle dependency: several flat deps, one slot + for leaf in d.deps: + t.add_dep(leaf) + t._slot_map.append((len(d.deps), list(d.names), list(d.acquired))) elif isinstance(d, exec_place): if exec_place_set: raise ValueError("Only one exec_place can be given") diff --git a/python/cuda_stf/tests/stf/test_bundles.py b/python/cuda_stf/tests/stf/test_bundles.py index 23a5f44cb47f..00d3ece8a3f2 100644 --- a/python/cuda_stf/tests/stf/test_bundles.py +++ b/python/cuda_stf/tests/stf/test_bundles.py @@ -107,3 +107,24 @@ def test_bundle_device_logical_data(): with ctx.task(B.rw()) as t: assert t.get(0).vals.__cuda_array_interface__["shape"] == (8,) ctx.finalize() + + +def test_stackable_bundle_tokens(): + """Bundles over stackable contexts: token bundles + launchable replay.""" + ctx = stf.stackable_context() + B = ctx.bundle(rigid=ctx.token(), soft=ctx.token()) + lx = ctx.logical_data(np.zeros(8, dtype=np.float64)) + + # whole-bundle modes distribute over the token fields; one slot each + with ctx.task(B.rw(), lx.rw()) as t: + # tokens are ordering-only: the bundle slot exists, its views are CAIs + # of the token deps' (empty) payloads; the plain dep is slot 1 + assert t.get(1).__cuda_array_interface__["shape"] == (8,) + + # per-field spelling with exclusion works identically to plain contexts + bd = B.dep(rigid=stf.AccessMode.READ, soft=stf.AccessMode.NONE) + assert bd.acquired == [True, False] + with ctx.task(bd, lx.read()): + pass + + ctx.finalize()