[STF] Add bundles: grouped logical data as a single task dependency

cudax/examples/stf/graph_algorithms/pagerank.cu

@@ -67,8 +67,14 @@ int main()
   std::vector<float> page_rank(num_vertices, init_rank);
   std::vector<float> 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<field<slice<int>, constant>, field<slice<int>, 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<scalar_view<float>>());
@@ -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<float>{}))
-        ->*[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<float>{}))
+        ->*[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]);
             };

cudax/examples/stf/linear_algebra/cg_csr.cu

@@ -21,20 +21,30 @@ using vector_t  = logical_data<slice<double>>;
 using scalar_t  = logical_data<scalar_view<double>>;
 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<slice<double>>, field<slice<size_t>, constant>, field<slice<size_t>, 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<slice<double>> val_handle;
-  mutable logical_data<slice<size_t>> row_handle;
-  mutable logical_data<slice<size_t>> 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

cudax/include/cuda/experimental/__stf/internal/backend_ctx.cuh

@@ -33,6 +33,7 @@
 #include <cuda/experimental/__places/machine.cuh>
 #include <cuda/experimental/__stf/allocators/block_allocator.cuh>
 #include <cuda/experimental/__stf/internal/async_resources_handle.cuh>
+#include <cuda/experimental/__stf/internal/bundle.cuh> // bundle-aware dependency overloads
 #include <cuda/experimental/__stf/internal/ctx_resource.cuh>
 #include <cuda/experimental/__stf/internal/execution_policy.cuh> // backend_ctx<T>::launch() uses execution_policy
 #include <cuda/experimental/__stf/internal/interpreted_execution_policy.cuh>
@@ -1151,7 +1152,8 @@ public:
   template <typename exec_place_t,
             typename S,
             typename... Deps,
-            typename = ::cuda::std::enable_if_t<::cuda::std::is_base_of_v<exec_place, exec_place_t>>>
+            typename = ::cuda::std::enable_if_t<::cuda::std::is_base_of_v<exec_place, exec_place_t>
+                                                && !reserved::any_bundle_dep_v<Deps...>>>
   auto parallel_for(exec_place_t e_place, S shape, Deps... deps)
   {
     if constexpr (::cuda::std::is_integral_v<S>)
@@ -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<exec_place, exec_place_t>>>
+            typename = ::cuda::std::enable_if_t<::cuda::std::is_base_of_v<exec_place, exec_place_t>
+                                                && !reserved::any_bundle_dep_v<Deps...>>>
   auto parallel_for(partitioner_t p, exec_place_t e_place, S shape, Deps... deps)
   {
     if constexpr (::cuda::std::is_integral_v<S>)
@@ -1183,12 +1186,125 @@ public:
     }
   }
 
-  template <typename S, typename... Deps>
+  template <typename S, typename... Deps, typename = ::cuda::std::enable_if_t<!reserved::any_bundle_dep_v<Deps...>>>
   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 <typename... Args, typename = ::cuda::std::enable_if_t<reserved::any_bundle_dep_v<Args...>>>
+  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 <typename... Args, typename = ::cuda::std::enable_if_t<reserved::any_bundle_dep_v<Args...>>>
+  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<exec_place, exec_place_t> && reserved::any_bundle_dep_v<Args...>,
+                             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 <typename S, typename... Args, ::cuda::std::enable_if_t<reserved::any_bundle_dep_v<Args...>, int> = 0>
+  auto parallel_for(S shape, Args... args)
+  {
+    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<exec_place, exec_place_t> && reserved::any_bundle_dep_v<Args...>,
+                             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 <typename thread_hierarchy_spec_t,
+            typename... Args,
+            typename = ::cuda::std::enable_if_t<reserved::any_bundle_dep_v<Args...>>>
+  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 <typename... Args, typename = ::cuda::std::enable_if_t<reserved::any_bundle_dep_v<Args...>>>
+  auto launch(exec_place e_place, Args... args)
+  {
+    return launch(par(par()), mv(e_place), mv(args)...);
+  }
+
+  template <typename... Args, typename = ::cuda::std::enable_if_t<reserved::any_bundle_dep_v<Args...>>>
+  auto launch(Args... args)
+  {
+    return launch(self().default_exec_place(), mv(args)...);
+  }
+
+  template <typename... Args, typename = ::cuda::std::enable_if_t<reserved::any_bundle_dep_v<Args...>>>
+  auto host_launch(Args... args)
+  {
+    return reserved::make_bundle_scope(
+      [&](auto... flat) {
+        return this->host_launch(mv(flat)...);
+      },
+      mv(args)...);
+  }
+
+  template <typename... Args, typename = ::cuda::std::enable_if_t<reserved::any_bundle_dep_v<Args...>>>
+  auto cuda_kernel(Args... args)
+  {
+    return reserved::make_bundle_scope(
+      [&](auto... flat) {
+        return this->cuda_kernel(mv(flat)...);
+      },
+      mv(args)...);
+  }
+
+  template <typename... Args, typename = ::cuda::std::enable_if_t<reserved::any_bundle_dep_v<Args...>>>
+  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()
   {