[INFRA] Import NVIDIA/CCCL upstream as optimization reference library

CCCL (CUDA C++ Core Libraries) provides:
- CUB: device/block/warp-level GPU primitives (reduce, scan, sort, topk)
- Thrust: high-level parallel algorithms (transform_reduce, sort, scan)
- libcudacxx: CUDA C++ standard library (atomics, barriers, memory)
- cudax: experimental features (memory resources, allocators)
- Tuning policies: per-SM hardware-specific algorithm parameters

Competition optimization vectors mapped to CCCL:
- Output TPS (83% weight): warp_reduce, block_reduce, device_topk
- Input TPS (14% weight): device_scan, block_load, prefetch
- Cache TPS (3% weight): prefix caching strategy patterns
- Memory (0.9 util): pooled/cached/buddy allocators

Source: https://github.com/NVIDIA/cccl (shallow clone, HEAD only)
License: Apache-2.0

cccl_upstream/cudax/test/stf/CMakeLists.txt

@@ -0,0 +1,336 @@
+set(
+  stf_test_sources
+  allocators/buddy_allocator.cu
+  context/context_resource_test.cu
+  context/simple_context_resource_test.cu
+  cpp/concurrency_test.cu
+  cpp/redundant_data.cu
+  cpp/redundant_data_different_modes.cu
+  cpp/scoped_graph_task.cu
+  cpp/task_get_stream.cu
+  cpp/test_pick_stream.cu
+  cpp/test_pick_stream_green_context.cu
+  cpp/user_streams.cu
+  cpp/wait_token.cu
+  cuda-samples/3_CUDA_Features/graphConditionalNodes/graphConditionalNodes.cu
+  dot/basic.cu
+  dot/graph_print_to_dot.cu
+  dot/sections.cu
+  dot/sections_2.cu
+  dot/sections_stackable.cu
+  dot/section_movable.cu
+  dot/with_events.cu
+  error_checks/ctx_mismatch.cu
+  error_checks/data_interface_mismatch.cu
+  error_checks/double_finalize.cu
+  error_checks/erase_frozen.cu
+  error_checks/misformed_tasks_dbl_end.cu
+  error_checks/misformed_tasks_dbl_start.cu
+  error_checks/non_managed_data.cu
+  error_checks/stackable_fence_in_nested.cu
+  error_checks/stackable_illegal_export.cu
+  error_checks/stackable_inconsistent_access_modes.cu
+  error_checks/stackable_wait_in_nested.cu
+  error_checks/uninitialized_data.cu
+  error_checks/write_frozen.cu
+  examples/cuda_kernels_driver.cu
+  examples/05-stencil-no-copy.cu
+  examples/05-stencil-places.cu
+  examples/05-stencil.cu
+  examples/05-stencil2d-places.cu
+  fhe/parse_arctyrex.cu
+  gnu/include_only.cpp
+  graph/concurrency_test.cu
+  graph/explicit_graph.cu
+  graph/explicit_graph_async.cu
+  graph/explicit_graph_while.cu
+  graph/explicit_graph_while-kernels.cu
+  graph/get_cache_stats.cu
+  graph/graph_cache_policy.cu
+  graph/graph_child_move_ownership.cu
+  graph/graph_ctx_low_level.cu
+  graph/static_graph_ctx.cu
+  hashtable/test.cu
+  interface/cuda_kernel_chain-add_deps.cu
+  interface/cuda_kernel_chain-add_deps_low_level.cu
+  interface/cuda_kernel_empty_args.cu
+  interface/move_operator.cu
+  local_stf/legacy_to_stf.cu
+  local_stf/legacy_to_stf_in_capture.cu
+  local_stf/logical_data_t_template.cu
+  local_stf/stackable.cu
+  local_stf/stackable2.cu
+  local_stf/stackable_add_deps.cu
+  local_stf/stackable_export.cu
+  local_stf/stackable_export2.cu
+  local_stf/stackable_move.cu
+  local_stf/stackable_nested_lib.cu
+  local_stf/stackable_nested_repeat.cu
+  local_stf/stackable_nested_while.cu
+  local_stf/stackable_nested.cu
+  local_stf/stackable_node_pool_growth.cu
+  local_stf/stream_ctx_lifetime_btb.cu
+  local_stf/stackable_read_only.cu
+  local_stf/stackable_threads.cu
+  local_stf/stackable_token.cu
+  local_stf/stackable_write_back.cu
+  local_stf/stackable_redundant_deps.cu
+  local_stf/stackable_tmp.cu
+  local_stf/task_dep_copy_move_semantics.cu
+  local_stf/threads_multiple_graphs.cu
+  places/affinity_gc.cu
+  places/cute_parallel_for.cu
+  places/managed.cu
+  places/managed_from_user.cu
+  places/non_current_device.cu
+  places/place_partition.cu
+  places/recursion.cu
+  places/execution_policy_kernel_launch_test.cu
+  reclaiming/graph.cu
+  reclaiming/graph_2.cu
+  reclaiming/graph_real_oom.cu
+  reclaiming/stream.cu
+  reductions/many_inc.cu
+  reductions/redux_test.cu
+  reductions/redux_test2.cu
+  reductions/slice2d_reduction.cu
+  reductions/slice_custom_op.cu
+  reductions/successive_reductions.cu
+  reductions/sum.cu
+  reductions/sum_array.cu
+  reductions/sum_multiple_places_no_refvalue.cu
+  slice/pinning.cu
+  stackable/graph_scope_test.cu
+  stencil/stencil-1D.cu
+  stress/empty_tasks.cu
+  stress/empty_tasks_alloc.cu
+  stress/kernel_chain.cu
+  stress/kernel_chain_fused.cu
+  stress/many_read.cu
+  stress/task_bench.cu
+  threads/axpy-threads-2.cu
+  # threads/axpy-threads.cu
+  utility/timing_with_fences.cu
+  utility/source_location_map.cu
+)
+
+set(
+  stf_test_codegen_sources
+  # algorithm/algorithm_with_read.cu
+  # algorithm/graph_algorithms.cu
+  # algorithm/in_graph_ctx.cu
+  # algorithm/nested.cu
+  allocators/adapter.cu
+  allocators/cap_tmp_buffers.cu
+  cpp/read_const.cu
+  cpp/reuse_computation.cu
+  cpp/reuse_computation_2.cu
+  error_checks/slice_check_bounds.cu
+  error_checks/unsatisfiable_spec.cu
+  examples/01-axpy-launch-ranges-cg.cu
+  examples/01-axpy-places.cu
+  # examples/09-nbody-algorithm.cu
+  examples/09-nbody-blocked.cu
+  examples/09-nbody.cu
+  freeze/constant_logical_data.cu
+  freeze/freeze.cu
+  freeze/freeze_rw.cu
+  freeze/freeze_untyped_rw.cu
+  freeze/freeze_write_back.cu
+  freeze/task_fence.cu
+  freeze/token.cu
+  graph/epoch.cu
+  graph/freeze_for_graph.cu
+  graph/freeze_nested_graphs.cu
+  graph/freeze_while_graphs.cu
+  graph/freeze_while_graphs_2.cu
+  graph/freeze_while_graphs_3.cu
+  # graph/graph_composition.cu
+  graph/graph_multiple_launch_resources.cu
+  graph/graph_tmp_data.cu
+  graph/many.cu
+  graph/multiple_graph_ctx.cu
+  green_context/axpy_gc.cu
+  green_context/cuda_graph.cu
+  green_context/gc_grid.cu
+  hash/ctx_hash.cu
+  hash/logical_data.cu
+  hashtable/fusion.cu
+  hashtable/fusion_reduction.cu
+  hashtable/parallel_for.cu
+  hashtable/parallel_for_shape.cu
+  interface/data_from_device.cu
+  interface/data_from_device_2.cu
+  interface/data_from_device_wb.cu
+  interface/graph_use_device_data.cu
+  interface/host_launch_deps.cu
+  interface/mix_stream_and_graph.cu
+  interface/mix_stream_and_graph_2.cu
+  interface/scal.cu
+  interface/scalar_div.cu
+  interface/scalar_interface.cu
+  interface/stream_add_callback.cu
+  local_stf/interop_cuda.cu
+  loop_dispatch/dispatch_on_streams.cu
+  loop_dispatch/loop_dispatch.cu
+  loop_dispatch/nested_loop_dispatch.cu
+  # loop_dispatch/stackable_loop_dispatch.cu # Disabled: double free or corruption (under investigation)
+  parallel_for/empty_shape_reduce.cu
+  parallel_for/fdtd.cu
+  parallel_for/parallel_for_all_devs.cu
+  parallel_for/parallel_for_box.cu
+  parallel_for/parallel_for_repeat.cu
+  parallel_for/test2_parallel_for_context.cu
+  parallel_for/tiled_loops.cu
+  parallel_for/parallel_for_host.cu
+  places/cuda_stream_place.cu
+  places/managed_from_shape.cu
+  reductions/reduce_sum.cu
+  reductions/successive_reductions_pfor.cu
+  reductions/sum_multiple_places.cu
+  reductions/write_back_after_redux.cu
+  stress/launch_overhead.cu
+  stress/launch_vs_parallelfor.cu
+  stress/parallel_for_overhead.cu
+  # threads/axpy-threads-pfor.cu   # Currently has a difficult-to-reproduce concurrency problem
+  threads/axpy-threads-graph.cu
+  threads/axpy-threads-graph-capture.cu
+)
+
+# Examples using CUBLAS, CUSOLVER...
+set(
+  stf_test_mathlib_sources
+  cuda-samples/0_Introduction/vectorAdd/vectorAdd_cudastf.cu
+  # Reduce compilation time by not adding this (useless) example
+  # cuda-samples/0_Introduction/vectorAdd/vectorAdd
+  cuda-samples/3_CUDA_Features/jacobiCudaGraphs/jacobi.cu
+  cuda-samples/3_CUDA_Features/jacobiCudaGraphs/jacobi_cudastf.cu
+  cuda-samples/4_CUDA_Libraries/conjugateGradientMultiDeviceCG/conjugateGradientMultiDeviceCG_custf.cu
+  cuda-samples/5_Domain_Specific/MonteCarloMultiGPU_cudastf/MonteCarloMultiGPU.cu
+  examples/07-cholesky-redux.cu
+  examples/07-cholesky-unified.cu
+  gnu/06-pdgemm.cpp
+  gnu/07-cholesky.cpp
+)
+
+set(
+  stf_unittested_headers
+  cuda/experimental/__stf/allocators/buddy_allocator.cuh
+  cuda/experimental/__stf/graph/graph_ctx.cuh
+  cuda/experimental/__stf/internal/async_resources_handle.cuh
+  cuda/experimental/__stf/internal/context.cuh
+  cuda/experimental/__stf/internal/execution_policy.cuh
+  cuda/experimental/__stf/internal/interpreted_execution_policy.cuh
+  cuda/experimental/__stf/internal/slice.cuh
+  cuda/experimental/__stf/internal/thread_hierarchy.cuh
+  cuda/experimental/__stf/internal/inner_shape.cuh
+  cuda/experimental/__stf/stream/stream_ctx.cuh
+  cuda/experimental/__stf/utility/cartesian_iterator.cuh
+  cuda/experimental/__stf/utility/cuda_safe_call.cuh
+  cuda/experimental/__stf/utility/dimensions.cuh
+  cuda/experimental/__stf/utility/handle.cuh
+  cuda/experimental/__stf/utility/hash.cuh
+  cuda/experimental/__stf/utility/memory.cuh
+  cuda/experimental/__stf/utility/scope_guard.cuh
+  cuda/experimental/__stf/stackable/stackable_ctx.cuh
+  cuda/experimental/__stf/utility/unittest.cuh
+)
+
+cccl_get_cudatoolkit()
+
+## cudax_add_stf_test
+#
+# Add an stf test executable and register it with ctest.
+#
+# target_name_var: Variable name to overwrite with the name of the test
+#   target. Useful for adding target information after creation.
+# source: The source file for the test.
+#
+# Additional args are passed to cudax_stf_configure_target.
+function(cudax_add_stf_test target_name_var source)
+  get_filename_component(dir ${source} DIRECTORY)
+  get_filename_component(filename ${source} NAME_WE)
+  if (dir)
+    set(filename "${dir}/${filename}")
+  endif()
+  string(REPLACE "/" "." test_name "${filename}")
+
+  set(test_target cudax.test.stf.${test_name})
+
+  cccl_add_executable(${test_target} SOURCES ${source} ADD_CTEST)
+  cudax_stf_configure_target(${test_target} ${ARGN})
+  target_link_libraries(${test_target} PRIVATE cudax.compiler_interface)
+
+  set(${target_name_var} ${test_target} PARENT_SCOPE)
+endfunction()
+
+## cudax_add_stf_unittest_header
+#
+# Add an stf unittested header executable and register it with ctest.
+#
+# Unittested headers contain a set of tests that are enabled by including
+# `unittest.cuh` and defining `UNITTESTED_FILE`.
+#
+# target_name_var: Variable name to overwrite with the name of the test
+#   target. Useful for adding target information after creation.
+# source: The source file for the test.
+#
+# Additional args are passed to cudax_stf_configure_target.
+function(cudax_add_stf_unittest_header target_name_var source)
+  get_filename_component(relative_path ${source} DIRECTORY)
+  get_filename_component(filename ${source} NAME_WE)
+
+  string(
+    REPLACE
+    "cuda/experimental/"
+    ""
+    test_label
+    "${relative_path}/${filename}"
+  )
+  string(REPLACE "/" "." test_label "${test_label}")
+
+  set(test_target "cudax.test.stf.unittest_headers.${test_label}")
+
+  # Pass the full path to configure_file (this is configured from cudax/tests/stf/
+  get_filename_component(
+    source_full_path
+    ../../../cudax/include/${source}
+    ABSOLUTE
+  )
+  set(source ${source_full_path})
+
+  set(ut_template "${cudax_SOURCE_DIR}/cmake/stf_header_unittest.in.cu")
+  set(ut_source "${cudax_BINARY_DIR}/unittest_headers/${test_target}.cu")
+  configure_file(${ut_template} ${ut_source} @ONLY)
+
+  cccl_add_executable(${test_target} SOURCES ${ut_source} ADD_CTEST)
+  cudax_stf_configure_target(${test_target} ${ARGN})
+  target_link_libraries(${test_target} PRIVATE cudax.compiler_interface)
+
+  set(${target_name_var} ${test_target} PARENT_SCOPE)
+endfunction()
+
+# Basic tests:
+foreach (source IN LISTS stf_test_sources)
+  cudax_add_stf_test(test_target "${source}")
+endforeach()
+
+if (cudax_ENABLE_CUDASTF_CODE_GENERATION)
+  foreach (source IN LISTS stf_test_codegen_sources)
+    cudax_add_stf_test(test_target "${source}")
+  endforeach()
+endif()
+
+# Tests with mathlib deps:
+if (cudax_ENABLE_CUDASTF_MATHLIBS)
+  foreach (source IN LISTS stf_test_mathlib_sources)
+    cudax_add_stf_test(test_target "${source}" LINK_MATHLIBS)
+  endforeach()
+endif()
+
+# Unittested headers
+foreach (source IN LISTS stf_unittested_headers)
+  cudax_add_stf_unittest_header(test_target "${source}")
+endforeach()
+
+add_subdirectory(static_error_checks)

cccl_upstream/cudax/test/stf/algorithm/algorithm_with_read.cu

@@ -0,0 +1,58 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+template <typename T>
+void init(context& ctx, logical_data<T> l, int val)
+{
+  ctx.parallel_for(l.shape(), l.write())->*[=] __device__(size_t i, auto s) {
+    s(i) = val;
+  };
+}
+
+int main()
+{
+  context ctx;
+
+  auto a = ctx.logical_data<int>(10000000);
+  auto b = ctx.logical_data<int>(10000000);
+  auto c = ctx.logical_data<int>(10000000);
+  auto d = ctx.logical_data<int>(10000000);
+
+  init(ctx, a, 12);
+  init(ctx, b, 35);
+  init(ctx, c, 42);
+  init(ctx, d, 17);
+
+  /* a += 1; a += b; */
+  auto fn = [](context ctx, logical_data<slice<int>> a, logical_data<slice<int>> b) {
+    ctx.parallel_for(a.shape(), a.rw())->*[] __device__(size_t i, auto sa) {
+      sa(i) += 1;
+    };
+    ctx.parallel_for(a.shape(), a.rw(), b.read())->*[] __device__(size_t i, auto sa, auto sb) {
+      sa(i) += sb(i);
+    };
+  };
+
+  algorithm alg;
+
+  for (size_t i = 0; i < 100; i++)
+  {
+    alg.run_as_task(fn, ctx, a.rw(), b.read());
+    alg.run_as_task(fn, ctx, a.rw(), c.read());
+    alg.run_as_task(fn, ctx, c.rw(), d.read());
+    alg.run_as_task(fn, ctx, d.rw(), a.read());
+  }
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/algorithm/graph_algorithms.cu

@@ -0,0 +1,108 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+template <typename ctx_t, typename T, typename T2>
+void lib_call(ctx_t& ctx, logical_data<T> a, logical_data<T2> b)
+{
+  nvtx_range r("lib_call");
+  // b *= 2
+  // a = a + b
+  // b = a
+  ctx.parallel_for(b.shape(), b.rw())->*[] __device__(size_t i, auto sb) {
+    sb(i) *= 2;
+  };
+
+  ctx.parallel_for(a.shape(), a.rw(), b.read())->*[] __device__(size_t i, auto sa, auto sb) {
+    sa(i) += sb(i);
+  };
+
+  ctx.parallel_for(a.shape(), a.read(), b.write())->*[] __device__(size_t i, auto sa, auto sb) {
+    sb(i) = sa(i);
+  };
+}
+
+template <typename T>
+void init(context& ctx, logical_data<T> l, int val)
+{
+  ctx.parallel_for(l.shape(), l.write())->*[=] __device__(size_t i, auto s) {
+    s(i) = val;
+  };
+}
+
+int main()
+{
+  context ctx;
+
+  auto a = ctx.logical_data<int>(size_t(10000000));
+  auto b = ctx.logical_data<int>(size_t(10000000));
+  auto c = ctx.logical_data<int>(size_t(10000000));
+  auto d = ctx.logical_data<int>(size_t(10000000));
+
+  init(ctx, a, 12);
+  init(ctx, b, 35);
+  init(ctx, c, 42);
+  init(ctx, d, 42);
+
+  auto fn = [](context ctx, logical_data<slice<int>> a, logical_data<slice<int>> b) {
+    ctx.parallel_for(b.shape(), b.rw())->*[] __device__(size_t i, auto sb) {
+      sb(i) *= 2;
+    };
+    ctx.parallel_for(a.shape(), a.rw(), b.read())->*[] __device__(size_t i, auto sa, auto sb) {
+      sa(i) += sb(i);
+    };
+    ctx.parallel_for(a.shape(), a.read(), b.write())->*[] __device__(size_t i, auto sa, auto sb) {
+      sb(i) = sa(i);
+    };
+  };
+
+  algorithm alg;
+
+  {
+    nvtx_range r("run");
+    for (size_t i = 0; i < 100; i++)
+    {
+      ctx.task(a.rw(), b.rw())->*[&alg, &fn, &ctx](cudaStream_t stream, slice<int> sa, slice<int> sb) {
+        alg.run(fn, ctx, stream, sa, sb);
+      };
+
+      ctx.task(b.rw(), c.rw())->*[&alg, &fn, &ctx](cudaStream_t stream, slice<int> sb, slice<int> sc) {
+        alg.run(fn, ctx, stream, sb, sc);
+      };
+
+      ctx.task(c.rw(), d.rw())->*[&alg, &fn, &ctx](cudaStream_t stream, slice<int> sc, slice<int> sd) {
+        alg.run(fn, ctx, stream, sc, sd);
+      };
+
+      ctx.task(d.rw(), a.rw())->*[&alg, &fn, &ctx](cudaStream_t stream, slice<int> sd, slice<int> sa) {
+        alg.run(fn, ctx, stream, sd, sa);
+      };
+    }
+  }
+
+  {
+    nvtx_range r("run_as_task");
+    for (size_t i = 0; i < 100; i++)
+    {
+      alg.run_as_task(fn, ctx, a.rw(), b.rw());
+
+      alg.run_as_task(fn, ctx, a.rw(), c.rw());
+
+      alg.run_as_task(fn, ctx, c.rw(), d.rw());
+
+      alg.run_as_task(fn, ctx, d.rw(), a.rw());
+    }
+  }
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/algorithm/in_graph_ctx.cu

@@ -0,0 +1,68 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+template <typename context_t, typename T>
+void init(context_t& ctx, logical_data<T> l, int val)
+{
+  ctx.parallel_for(l.shape(), l.write())->*[=] __device__(size_t i, auto s) {
+    s(i) = val;
+  };
+}
+
+int main()
+{
+  //    context ctx = graph_ctx();
+  graph_ctx ctx;
+
+  auto a = ctx.logical_data<int>(size_t(1000000));
+  auto b = ctx.logical_data<int>(size_t(1000000));
+  auto c = ctx.logical_data<int>(size_t(1000000));
+  auto d = ctx.logical_data<int>(size_t(1000000));
+
+  init(ctx, a, 12);
+  init(ctx, b, 35);
+  init(ctx, c, 42);
+  init(ctx, d, 42);
+
+  auto fn = [](context ctx, logical_data<slice<int>> a, logical_data<slice<int>> b) {
+    ctx.parallel_for(a.shape(), a.rw())->*[] __device__(size_t i, auto sa) {
+      sa(i) *= 3;
+    };
+    ctx.parallel_for(b.shape(), b.rw())->*[] __device__(size_t i, auto sb) {
+      sb(i) *= 2;
+    };
+
+    ctx.parallel_for(a.shape(), a.rw(), b.read())->*[] __device__(size_t i, auto sa, auto sb) {
+      sa(i) += sb(i);
+    };
+
+    ctx.parallel_for(a.shape(), a.read(), b.write())->*[] __device__(size_t i, auto sa, auto sb) {
+      sb(i) = sa(i);
+    };
+  };
+
+  algorithm alg;
+
+  for (size_t i = 0; i < 5; i++)
+  {
+    alg.run_as_task(fn, ctx, a.rw(), b.rw());
+    alg.run_as_task(fn, ctx, a.rw(), c.rw());
+    alg.run_as_task(fn, ctx, c.rw(), d.rw());
+    alg.run_as_task(fn, ctx, d.rw(), a.rw());
+  }
+
+  ctx.finalize();
+
+  // cudaGraphDebugDotPrint(ctx.get_graph(), "pif.dot", 0);
+}

cccl_upstream/cudax/test/stf/algorithm/nested.cu

@@ -0,0 +1,67 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+template <typename T>
+void init(context& ctx, logical_data<T> l, int val)
+{
+  ctx.parallel_for(l.shape(), l.write())->*[=] __device__(size_t i, auto s) {
+    s(i) = val;
+  };
+}
+
+int main()
+{
+  context ctx;
+
+  auto a = ctx.logical_data<int>(size_t(1000000));
+  auto b = ctx.logical_data<int>(size_t(1000000));
+  auto c = ctx.logical_data<int>(size_t(1000000));
+  auto d = ctx.logical_data<int>(size_t(1000000));
+
+  init(ctx, a, 12);
+  init(ctx, b, 35);
+  init(ctx, c, 42);
+  init(ctx, d, 17);
+
+  auto fn1 = [](context ctx, logical_data<slice<int>> a) {
+    ctx.parallel_for(a.shape(), a.rw())->*[] __device__(size_t i, auto sa) {
+      sa(i) += 1;
+    };
+  };
+
+  algorithm alg1;
+
+  auto fn2 = [&alg1, &fn1](context ctx, logical_data<slice<int>> a, logical_data<slice<int>> b) {
+    alg1.run_as_task(fn1, ctx, a.rw());
+    alg1.run_as_task(fn1, ctx, b.rw());
+    ctx.parallel_for(a.shape(), a.rw(), b.read())->*[] __device__(size_t i, auto sa, auto sb) {
+      sa(i) += sb(i);
+    };
+    ctx.parallel_for(a.shape(), a.read(), b.write())->*[] __device__(size_t i, auto sa, auto sb) {
+      sb(i) = sa(i);
+    };
+  };
+
+  algorithm alg2;
+
+  for (size_t i = 0; i < 100; i++)
+  {
+    alg2.run_as_task(fn2, ctx, a.rw(), b.rw());
+    alg2.run_as_task(fn2, ctx, a.rw(), c.rw());
+    alg2.run_as_task(fn2, ctx, c.rw(), d.rw());
+    alg2.run_as_task(fn2, ctx, d.rw(), a.rw());
+  }
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/allocators/adapter.cu

@@ -0,0 +1,60 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/experimental/__stf/allocators/adapters.cuh>
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+int main()
+{
+  double* d_ptrA;
+  const size_t N     = 128 * 1024;
+  const size_t NITER = 10;
+
+  // User allocated memory
+  cuda_safe_call(cudaMalloc(&d_ptrA, N * sizeof(double)));
+
+  async_resources_handle handle;
+
+  cudaStream_t stream;
+  cuda_safe_call(cudaStreamCreate(&stream));
+
+  for (size_t i = 0; i < NITER; i++)
+  {
+    graph_ctx ctx(stream, handle);
+
+    // The uncached allocator of the context will be using cudaMallocAsync(...,
+    // stream) to avoid creating memory nodes in the graph (because they are
+    // costly and caching the graph also keeps memory allocated)
+    auto wrapper = stream_adapter(ctx, stream);
+
+    ctx.set_allocator(block_allocator<buddy_allocator>(ctx, wrapper.allocator()));
+
+    auto A = ctx.logical_data(make_slice(d_ptrA, N), data_place::current_device());
+
+    for (size_t k = 0; k < 4; k++)
+    {
+      auto tmp  = ctx.logical_data(A.shape());
+      auto tmp2 = ctx.logical_data(A.shape());
+      // Test device and managed memory
+      ctx.parallel_for(A.shape(), A.read(), tmp.write(), tmp2.write(data_place::managed()))
+          ->*[] __device__(size_t i, auto a, auto tmp, auto tmp2) {
+                tmp(i)  = a(i);
+                tmp2(i) = a(i);
+              };
+    }
+
+    ctx.finalize();
+
+    wrapper.clear();
+  }
+  cuda_safe_call(cudaStreamSynchronize(stream));
+}

cccl_upstream/cudax/test/stf/allocators/buddy_allocator.cu

@@ -0,0 +1,43 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/experimental/stf.cuh>
+
+/**
+ * @brief Ensure the buddy allocation is working properly on the different backends
+ */
+
+using namespace cuda::experimental::stf;
+
+template <typename ctx_t>
+void test_buddy()
+{
+  ctx_t ctx;
+  ctx.set_allocator(block_allocator<buddy_allocator>(ctx));
+
+  std::vector<logical_data<slice<char>>> data;
+  for (size_t i = 0; i < 10; i++)
+  {
+    size_t s = (1 + i % 8) * 1024ULL * 1024ULL;
+    auto l   = ctx.logical_data(shape_of<slice<char>>(s));
+    data.push_back(l);
+
+    ctx.task(l.write())->*[](cudaStream_t, auto) {};
+  }
+
+  ctx.finalize();
+}
+
+int main(int, char**)
+{
+  test_buddy<stream_ctx>();
+  test_buddy<graph_ctx>();
+  test_buddy<context>();
+}

cccl_upstream/cudax/test/stf/allocators/cap_tmp_buffers.cu

@@ -0,0 +1,56 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+int main(int, char**)
+{
+  context ctx;
+
+  const size_t PART_SIZE = 1024;
+  const size_t PART_CNT  = 64;
+
+  pooled_allocator_config config;
+  config.max_entries_per_place = 8;
+  auto fixed_alloc             = block_allocator<pooled_allocator>(ctx, config);
+
+  /* Create a large device buffer which will be used part by part. */
+  double* dA;
+  cuda_safe_call(cudaMalloc(&dA, PART_SIZE * PART_CNT * sizeof(double)));
+
+  for (size_t p = 0; p < PART_CNT; p++)
+  {
+    /* Create a logical data from a subset of the existing device buffer */
+    auto Ap = ctx.logical_data(make_slice(&dA[p * PART_SIZE], PART_SIZE), data_place::current_device());
+
+    ctx.parallel_for(Ap.shape(), Ap.write()).set_symbol("init_Ap")->*[p, PART_SIZE] __device__(size_t i, auto ap) {
+      ap(i) = 1.0 * (i + p * PART_SIZE);
+    };
+
+    auto tmp = ctx.logical_data(Ap.shape());
+    tmp.set_allocator(fixed_alloc);
+
+    ctx.parallel_for(Ap.shape(), Ap.read(), tmp.write()).set_symbol("set_tmp")->*
+      [] __device__(size_t i, auto ap, auto tmp) {
+        tmp(i) = 2.0 * ap(i);
+      };
+
+    ctx.parallel_for(Ap.shape(), Ap.write(), tmp.read()).set_symbol("update_Ap")
+        ->*[] __device__(size_t i, auto ap, auto tmp) {
+              ap(i) = tmp(i);
+            };
+  }
+
+  ctx.finalize();
+
+  cuda_safe_call(cudaFree(dA));
+}

cccl_upstream/cudax/test/stf/context/context_resource_test.cu

@@ -0,0 +1,222 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2025 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+//! \file
+//! \brief Test ctx_resource management with different context types
+
+#include <cuda/experimental/__stf/internal/context.cuh>
+
+#include <atomic>
+#include <memory>
+
+using namespace cuda::experimental::stf;
+
+namespace
+{
+// Counters for tracking resource lifecycle
+std::atomic<int> stream_resource_construct_count{0};
+std::atomic<int> stream_resource_release_count{0};
+std::atomic<int> callback_resource_construct_count{0};
+std::atomic<int> callback_resource_release_count{0};
+
+// Test resource that requires a stream for release
+class test_stream_resource : public ctx_resource
+{
+public:
+  test_stream_resource()
+  {
+    stream_resource_construct_count.fetch_add(1);
+  }
+
+  ~test_stream_resource() override = default;
+
+  void release(cudaStream_t stream) noexcept override
+  {
+    // Simulate async resource release that needs a stream
+    cudaEvent_t event;
+    cuda_safe_call(cudaEventCreate(&event));
+    cuda_safe_call(cudaEventRecord(event, stream));
+    cuda_safe_call(cudaEventSynchronize(event)); // Wait for completion
+    cuda_safe_call(cudaEventDestroy(event));
+
+    stream_resource_release_count.fetch_add(1);
+  }
+
+  bool can_release_in_callback() const noexcept override
+  {
+    return false; // This resource needs a stream
+  }
+};
+
+// Test resource that can be released in a host callback
+class test_callback_resource : public ctx_resource
+{
+public:
+  test_callback_resource()
+  {
+    callback_resource_construct_count.fetch_add(1);
+  }
+
+  ~test_callback_resource() override = default;
+
+  void release(cudaStream_t /*stream*/) noexcept override
+  {
+    // Should not be called for callback resources
+    assert(false && "release() should not be called for callback resources");
+  }
+
+  bool can_release_in_callback() const noexcept override
+  {
+    return true; // This resource can be released in a callback
+  }
+
+  void release_in_callback() noexcept override
+  {
+    // Simulate host-side resource cleanup
+    callback_resource_release_count.fetch_add(1);
+  }
+};
+
+void reset_counters()
+{
+  stream_resource_construct_count.store(0);
+  stream_resource_release_count.store(0);
+  callback_resource_construct_count.store(0);
+  callback_resource_release_count.store(0);
+}
+
+void check_all_resources_released()
+{
+  // Verify all constructed resources were properly released
+  EXPECT(stream_resource_construct_count.load() == stream_resource_release_count.load());
+  EXPECT(callback_resource_construct_count.load() == callback_resource_release_count.load());
+}
+
+template <typename CtxType>
+void test_context_resources()
+{
+  reset_counters();
+
+  CtxType ctx;
+
+  // Add a simple host launch to ensure context has some work
+  ctx.host_launch()->*[]() {
+    // Trivial workload - just increment a counter
+    static std::atomic<int> work_counter{0};
+    work_counter.fetch_add(1);
+  };
+
+  // Add various types of resources
+  const int num_stream_resources   = 3;
+  const int num_callback_resources = 2;
+
+  // Add stream-dependent resources
+  for (int i = 0; i < num_stream_resources; ++i)
+  {
+    auto resource = ::std::make_shared<test_stream_resource>();
+    ctx.add_resource(resource);
+  }
+
+  // Add callback resources
+  for (int i = 0; i < num_callback_resources; ++i)
+  {
+    auto resource = ::std::make_shared<test_callback_resource>();
+    ctx.add_resource(resource);
+  }
+
+  // Verify resources were constructed
+  EXPECT(stream_resource_construct_count.load() == num_stream_resources);
+  EXPECT(callback_resource_construct_count.load() == num_callback_resources);
+  EXPECT(stream_resource_release_count.load() == 0); // Not released yet
+  EXPECT(callback_resource_release_count.load() == 0); // Not released yet
+
+  // Finalize the context - this should release resources automatically
+  ctx.finalize();
+
+  // Verify all resources were released
+  EXPECT(stream_resource_release_count.load() == num_stream_resources);
+  EXPECT(callback_resource_release_count.load() == num_callback_resources);
+
+  check_all_resources_released();
+}
+
+void test_graph_ctx_manual_resource_release()
+{
+  reset_counters();
+
+  graph_ctx ctx;
+
+  // Add a simple host launch with work counter
+  std::atomic<int> work_counter{0};
+  ctx.host_launch()->*[&work_counter]() {
+    work_counter.fetch_add(1);
+  };
+
+  // Add resources
+  const int num_resources = 2;
+  for (int i = 0; i < num_resources; ++i)
+  {
+    ctx.add_resource(std::make_shared<test_stream_resource>());
+    ctx.add_resource(std::make_shared<test_callback_resource>());
+  }
+
+  EXPECT(stream_resource_construct_count.load() == num_resources);
+  EXPECT(callback_resource_construct_count.load() == num_resources);
+
+  // Resources should not be released yet
+  EXPECT(stream_resource_release_count.load() == 0);
+  EXPECT(callback_resource_release_count.load() == 0);
+
+  // Generate the graph using finalize_as_graph
+  ::std::shared_ptr<cudaGraph_t> graph = ctx.finalize_as_graph();
+
+  // Create stream and instantiate graph for multiple launches
+  cudaStream_t test_stream;
+  cuda_safe_call(cudaStreamCreate(&test_stream));
+
+  cudaGraphExec_t graphExec;
+  cuda_safe_call(cudaGraphInstantiate(&graphExec, *graph, nullptr, nullptr, 0));
+
+  // Launch the graph multiple times
+  const int num_launches = 3;
+  for (int i = 0; i < num_launches; i++)
+  {
+    cuda_safe_call(cudaGraphLaunch(graphExec, test_stream));
+  }
+
+  // Manually release resources after graph executions
+  ctx.release_resources(test_stream);
+  cuda_safe_call(cudaStreamSynchronize(test_stream));
+
+  // Verify the work was executed
+  EXPECT(work_counter.load() == num_launches);
+
+  // Now resources should be released
+  EXPECT(stream_resource_release_count.load() == num_resources);
+  EXPECT(callback_resource_release_count.load() == num_resources);
+
+  // Clean up
+  cuda_safe_call(cudaGraphExecDestroy(graphExec));
+  cuda_safe_call(cudaStreamDestroy(test_stream));
+
+  check_all_resources_released();
+}
+} // anonymous namespace
+
+int main()
+{
+  // Test with different context types
+  test_context_resources<context>();
+  test_context_resources<stream_ctx>();
+  test_context_resources<graph_ctx>();
+
+  // Test manual resource release (graph_ctx only for the sake of simplicity)
+  test_graph_ctx_manual_resource_release();
+}

cccl_upstream/cudax/test/stf/context/simple_context_resource_test.cu

@@ -0,0 +1,114 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2025 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+//! \file
+//! \brief Simple test demonstrating ctx_resource management with generic context
+
+#include <cuda/experimental/__stf/internal/context.cuh>
+
+#include <atomic>
+#include <memory>
+
+using namespace cuda::experimental::stf;
+
+// Simple test resource that tracks its lifecycle
+class simple_test_resource : public ctx_resource
+{
+  static ::std::atomic<int> alive_count;
+
+public:
+  simple_test_resource()
+  {
+    alive_count.fetch_add(1);
+  }
+
+  ~simple_test_resource() override
+  {
+    alive_count.fetch_sub(1);
+  }
+
+  void release(cudaStream_t /*stream*/) noexcept override
+  {
+    // No special release action needed for this test
+  }
+
+  bool can_release_in_callback() const noexcept override
+  {
+    return true; // Can be released in a host callback
+  }
+
+  void release_in_callback() noexcept override
+  {
+    // Host-side cleanup - nothing to do for this simple test
+  }
+
+  static int get_alive_count()
+  {
+    return alive_count.load();
+  }
+};
+
+::std::atomic<int> simple_test_resource::alive_count{0};
+
+int main()
+{
+  // Test with generic context (defaults to stream_ctx)
+  {
+    context ctx; // Default initialization as stream_ctx
+
+    EXPECT(simple_test_resource::get_alive_count() == 0);
+
+    // Add a simple host launch with some work
+    ctx.host_launch()->*[]() {
+      // Trivial workload
+    };
+
+    // Add some resources to the context
+    for (int i = 0; i < 5; ++i)
+    {
+      auto resource = ::std::make_shared<simple_test_resource>();
+      ctx.add_resource(resource);
+    }
+
+    // Verify resources are alive
+    EXPECT(simple_test_resource::get_alive_count() == 5);
+
+    // Finalize the context - this should release all resources
+    ctx.finalize();
+  } // Context goes out of scope
+
+  // All resources should have been cleaned up
+  EXPECT(simple_test_resource::get_alive_count() == 0);
+
+  // Test with graph context through generic interface
+  {
+    context ctx = graph_ctx(); // Explicitly use graph backend
+
+    EXPECT(simple_test_resource::get_alive_count() == 0);
+
+    ctx.host_launch()->*[]() {
+      // Trivial workload
+    };
+
+    // Add resources
+    for (int i = 0; i < 3; ++i)
+    {
+      ctx.add_resource(::std::make_shared<simple_test_resource>());
+    }
+
+    EXPECT(simple_test_resource::get_alive_count() == 3);
+
+    ctx.finalize();
+  }
+
+  EXPECT(simple_test_resource::get_alive_count() == 0);
+
+  return 0;
+}

cccl_upstream/cudax/test/stf/cpp/concurrency_test.cu

@@ -0,0 +1,98 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/experimental/stf.cuh>
+
+/*
+ * The goal of this test is to ensure that using read access modes actually
+ * results in concurrent tasks
+ */
+
+using namespace cuda::experimental::stf;
+
+/**
+ * @brief Call `__nanosleep` (potentially repeatedly) to sleep `nanoseconds` nanoseconds. Supports sleep times longer
+ * than 4 billion nanoseconds (i.e. 4 seconds).
+ *
+ * @param nanoseconds how many nanoseconds to sleep
+ * @return void
+ */
+__global__ void nano_sleep(unsigned long long nanoseconds)
+{
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 700)
+  static constexpr auto m = std::numeric_limits<unsigned int>::max();
+  for (;;)
+  {
+    if (nanoseconds > m)
+    {
+      __nanosleep(m);
+      nanoseconds -= m;
+    }
+    else
+    {
+      __nanosleep(static_cast<unsigned int>(nanoseconds));
+      break;
+    }
+  }
+#else
+  const clock_t end = clock() + nanoseconds / (1000000000ULL / CLOCKS_PER_SEC);
+  while (clock() < end)
+  {
+    // busy wait
+  }
+#endif
+}
+
+void run(context& ctx, int NTASKS, int ms)
+{
+  int dummy[1];
+  auto handle = ctx.logical_data(dummy);
+
+  ctx.task().add_deps(handle.rw())->*[](cudaStream_t stream) {
+    nano_sleep<<<1, 1, 0, stream>>>(0);
+  };
+
+  for (int iter = 0; iter < 10; iter++)
+  {
+    for (int k = 0; k < NTASKS; k++)
+    {
+      ctx.task().add_deps(handle.read())->*[&](cudaStream_t stream) {
+        nano_sleep<<<1, 1, 0, stream>>>(ms * 1000ULL * 1000ULL);
+      };
+    }
+
+    ctx.task().add_deps(handle.rw())->*[&](cudaStream_t stream) {
+      nano_sleep<<<1, 1, 0, stream>>>(0);
+    };
+  }
+
+  ctx.finalize();
+}
+
+int main(int argc, char** argv)
+{
+  int NTASKS = 256;
+  int ms     = 40;
+
+  if (argc > 1)
+  {
+    NTASKS = atoi(argv[1]);
+  }
+
+  if (argc > 2)
+  {
+    ms = atoi(argv[2]);
+  }
+
+  context ctx;
+  run(ctx, NTASKS, ms);
+  ctx = graph_ctx();
+  run(ctx, NTASKS, ms);
+}

cccl_upstream/cudax/test/stf/cpp/read_const.cu

@@ -0,0 +1,93 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief This test ensures that read() can be called on const logical_data
+ *        (typed and untyped).
+ */
+
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+struct foo
+{
+  // Intentionally choose an odd (actually prime) size
+  foo(context& ctx)
+  {
+    l = ctx.logical_data(shape_of<slice<int>>(50867));
+  }
+
+  void set(context& ctx, int val)
+  {
+    ctx.parallel_for(l.shape(), l.write())->*[=] _CCCL_DEVICE(size_t i, auto dl) {
+      dl(i) = val;
+    };
+  }
+
+  void copy_from(context& ctx, const foo& other)
+  {
+    ctx.parallel_for(l.shape(), l.write(), other.l.read())->*[=] _CCCL_DEVICE(size_t i, auto dl, auto dotherl) {
+      dl(i) = dotherl(i);
+    };
+  }
+
+  void ensure(context& ctx, int val)
+  {
+    std::ignore = val;
+    ctx.parallel_for(l.shape(), l.read())->*[=] _CCCL_DEVICE(size_t i, auto dl) {
+      assert(dl(i) == val);
+    };
+  }
+
+  auto& get_l() const
+  {
+    return l;
+  }
+
+  logical_data<slice<int>> l;
+};
+
+void read_only_access(context& ctx, const foo& f)
+{
+  ctx.parallel_for(f.l.shape(), f.l.read())->*[] _CCCL_DEVICE(size_t i, auto dl) {
+    // no-op
+  };
+
+  ctx.parallel_for(f.get_l().shape(), f.get_l().read())->*[] _CCCL_DEVICE(size_t i, auto dl) {
+    // no-op
+  };
+}
+
+void read_only_access_untyped(const logical_data_untyped& ld)
+{
+  auto dep = ld.read();
+  (void) dep;
+}
+
+int main()
+{
+  context ctx;
+
+  foo A(ctx);
+  A.set(ctx, 42);
+  A.ensure(ctx, 42);
+  foo B(ctx);
+  B.copy_from(ctx, A);
+  B.ensure(ctx, 42);
+
+  read_only_access(ctx, A);
+
+  const logical_data_untyped& ld_untyped = A.l;
+  read_only_access_untyped(ld_untyped);
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/cpp/redundant_data.cu

@@ -0,0 +1,86 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief Ensure we can use the same logical data multiple time in a task
+ */
+
+#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+template <typename T>
+__global__ void diff_cnt(int n, T* x, T* y, int* delta)
+{
+  int tid      = blockIdx.x * blockDim.x + threadIdx.x;
+  int nthreads = gridDim.x * blockDim.x;
+
+  for (int ind = tid; ind < n; ind += nthreads)
+  {
+    if (y[ind] != x[ind])
+    {
+      atomicAdd(delta, 1);
+    }
+  }
+}
+
+template <typename Ctx, typename T>
+void compare_two_vectors(Ctx& ctx, logical_data<T>& a, logical_data<T>& b, int& delta)
+{
+  auto delta_cnt = ctx.logical_data(make_slice(&delta, 1));
+  const auto n   = a.shape().extent(0);
+
+  // Count the number of differences
+  ctx.task(a.read(), b.read(), delta_cnt.rw())->*[=](cudaStream_t stream, auto da, auto db, auto ddelta) {
+    diff_cnt<<<16, 128, 0, stream>>>(static_cast<int>(n), da.data_handle(), db.data_handle(), ddelta.data_handle());
+  };
+
+  // Read that value on the host
+  ctx.host_launch(delta_cnt.read())->*[&](auto /*unused*/) {};
+}
+
+static const size_t N = 12;
+
+template <class Ctx>
+void run(double (&X)[N], double (&Y)[N])
+{
+  Ctx ctx;
+  auto handle_X = ctx.logical_data(X);
+  auto handle_Y = ctx.logical_data(Y);
+
+  int ret1 = 0, ret2 = 0;
+
+  compare_two_vectors(ctx, handle_X, handle_Y, ret1);
+  compare_two_vectors(ctx, handle_X, handle_X, ret2);
+
+  ctx.finalize();
+
+  // After sync, we can inspect the returned values.
+  // First two vectors are different
+  assert(ret1 > 0);
+  // Other two vectors are equal
+  assert(ret2 == 0);
+}
+
+int main()
+{
+  double X[N], Y[N];
+
+  for (size_t ind = 0; ind < N; ind++)
+  {
+    X[ind] = 1.0 * ind;
+    Y[ind] = 2.0 * ind - 3.0;
+  }
+
+  run<stream_ctx>(X, Y);
+  run<graph_ctx>(X, Y);
+}

cccl_upstream/cudax/test/stf/cpp/redundant_data_different_modes.cu

@@ -0,0 +1,55 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief Ensure we can use the same logical data multiple times in the same
+ *        task even with different access modes (which should be combined)
+ */
+
+#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+// a = b + 1;
+template <typename T>
+__global__ void add(T* a, const T* b)
+{
+  *a = *b + 1;
+}
+
+template <class Ctx>
+void run()
+{
+  Ctx ctx;
+
+  int var         = 42;
+  auto var_handle = ctx.logical_data(make_slice(&var, 1));
+
+  // da and db are for the same variable : we expect it to be equivalent to a RW access
+  ctx.task(var_handle.write(), var_handle.read())->*[](cudaStream_t stream, auto da, auto db) {
+    add<<<1, 1, 0, stream>>>(da.data_handle(), db.data_handle());
+  };
+
+  // Read that value on the host
+  ctx.host_launch(var_handle.read())->*[](auto da) {
+    [[maybe_unused]] int result = *da.data_handle();
+    assert(result == 43);
+  };
+
+  ctx.finalize();
+}
+
+int main()
+{
+  run<stream_ctx>();
+  run<graph_ctx>();
+}

cccl_upstream/cudax/test/stf/cpp/reuse_computation.cu

@@ -0,0 +1,80 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/experimental/__stf/utility/run_once.cuh>
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+int main()
+{
+  context ctx;
+
+  const int N  = 16;
+  size_t niter = 12;
+
+  int A[N];
+
+  for (int i = 0; i < N; i++)
+  {
+    A[i] = 2 * i + 1;
+  }
+
+  auto lres = ctx.logical_data(A);
+
+  for (size_t k = 0; k < niter; k++)
+  {
+    auto ltmp = ctx.logical_data(lres.shape());
+    ctx.parallel_for(ltmp.shape(), ltmp.write())->*[] __device__(size_t i, auto tmp) {
+      tmp(i) = i;
+    };
+
+    ctx.parallel_for(lres.shape(), ltmp.read(), lres.rw())->*[] __device__(size_t i, auto tmp, auto res) {
+      res(i) += tmp(i);
+    };
+  }
+
+  for (size_t k = 0; k < niter; k++)
+  {
+    auto ltmp = run_once()->*[&]() {
+      // Ensure this is only done once !
+      static bool done = false;
+      EXPECT(!done);
+      done = true;
+
+      auto ltmp = ctx.logical_data(lres.shape());
+      ctx.parallel_for(ltmp.shape(), ltmp.write())->*[] __device__(size_t i, auto tmp) {
+        tmp(i) = i;
+      };
+      return ltmp;
+    };
+
+    auto ltmp2 = run_once(size_t(k % 4))->*[&](size_t val) {
+      // fprintf(stderr, "COMPUTE FOR %ld\n", val);
+
+      auto ltmp = ctx.logical_data(lres.shape());
+      ctx.parallel_for(ltmp.shape(), ltmp.write())->*[val] __device__(size_t i, auto tmp) {
+        tmp(i) = val;
+      };
+      return ltmp;
+    };
+
+    ctx.parallel_for(lres.shape(), ltmp.read(), lres.rw())->*[] __device__(size_t i, auto tmp, auto res) {
+      res(i) += tmp(i);
+    };
+  }
+
+  ctx.finalize();
+
+  for (int i = 0; i < N; i++)
+  {
+    EXPECT(A[i] == (2 * i + 1) + 2 * i * niter);
+  }
+}

cccl_upstream/cudax/test/stf/cpp/reuse_computation_2.cu

@@ -0,0 +1,65 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/experimental/__stf/utility/run_once.cuh>
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+int main()
+{
+  context ctx;
+
+  const int N  = 16;
+  size_t niter = 12;
+
+  int A[N];
+
+  for (int i = 0; i < N; i++)
+  {
+    A[i] = 2 * i + 1;
+  }
+
+  auto lres = ctx.logical_data(A);
+
+  for (size_t k = 0; k < niter; k++)
+  {
+    auto ltmp = ctx.logical_data(lres.shape());
+    ctx.parallel_for(ltmp.shape(), ltmp.write())->*[k] __device__(size_t i, auto tmp) {
+      tmp(i) = (k % 2) * i;
+    };
+
+    ctx.parallel_for(lres.shape(), ltmp.read(), lres.rw())->*[] __device__(size_t i, auto tmp, auto res) {
+      res(i) += tmp(i);
+    };
+  }
+
+  for (size_t k = 0; k < niter; k++)
+  {
+    auto ltmp = run_once(k)->*[&](size_t k) {
+      auto out = ctx.logical_data(lres.shape());
+      ctx.parallel_for(out.shape(), out.write())->*[k] __device__(size_t i, auto tmp) {
+        tmp(i) = (k % 2) * i;
+      };
+      return out;
+    };
+
+    ctx.parallel_for(lres.shape(), ltmp.read(), lres.rw())->*[] __device__(size_t i, auto tmp, auto res) {
+      res(i) += tmp(i);
+    };
+  }
+
+  ctx.finalize();
+
+  for (int i = 0; i < N; i++)
+  {
+    EXPECT(A[i] == (2 * i + 1) + 2 * i * niter / 2);
+  }
+}

cccl_upstream/cudax/test/stf/cpp/scoped_graph_task.cu

@@ -0,0 +1,103 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief Show how we can create tasks in the CUDA graph backend by using the
+ *        actual CUDA graph API in tasks (instead of relying on graph capture
+ *        implicitly)
+ */
+
+#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
+
+#include <iostream>
+
+using namespace cuda::experimental::stf;
+
+int main(int argc, char** argv)
+{
+  const size_t n = 12;
+
+  double X[n];
+  double Y[n];
+
+  for (size_t ind = 0; ind < n; ind++)
+  {
+    X[ind] = 1.0 * ind + 42;
+    Y[ind] = 0.0;
+  }
+
+  // We here do not assume there is a valid copy on the host and only provide
+  // constant parameters
+  graph_ctx ctx;
+  auto handle_X = ctx.logical_data(X);
+  handle_X.set_symbol("x");
+  auto handle_Y = ctx.logical_data(Y);
+  handle_Y.set_symbol("y");
+  auto handle_TMP = ctx.logical_data<double>(n);
+  handle_TMP.set_symbol("tmp");
+
+  int NITER = 4;
+  for (int iter = 0; iter < NITER; iter++)
+  {
+    // We swap X and Y using TMP as temporary buffer
+    // TMP = X
+    // X = Y
+    // Y = TMP
+    ctx.task(exec_place::current_device(), handle_X.rw(), handle_Y.rw(), handle_TMP.write())
+        ->*[&](cudaGraph_t child_graph, auto d_x, auto d_y, auto d_tmp) {
+              // TMP = X
+              cudaGraphNode_t cpy_tmp_to_x;
+              cuda_try(cudaGraphAddMemcpyNode1D(
+                &cpy_tmp_to_x,
+                child_graph,
+                nullptr,
+                0,
+                d_tmp.data_handle(),
+                d_x.data_handle(),
+                n * sizeof(double),
+                cudaMemcpyDeviceToDevice));
+
+              // X = Y
+              cudaGraphNode_t cpy_x_to_y;
+              cuda_try(cudaGraphAddMemcpyNode1D(
+                &cpy_x_to_y,
+                child_graph,
+                &cpy_tmp_to_x,
+                1,
+                d_x.data_handle(),
+                d_y.data_handle(),
+                n * sizeof(double),
+                cudaMemcpyDeviceToDevice));
+
+              // Y = TMP
+              cudaGraphNode_t cpy_tmp_to_y;
+              cuda_try(cudaGraphAddMemcpyNode1D(
+                &cpy_tmp_to_y,
+                child_graph,
+                &cpy_x_to_y,
+                1,
+                d_y.data_handle(),
+                d_tmp.data_handle(),
+                n * sizeof(double),
+                cudaMemcpyDeviceToDevice));
+            };
+  }
+
+  ctx.submit();
+
+  if (argc > 1)
+  {
+    std::cout << "Generating DOT output in " << argv[1] << '\n';
+    ctx.print_to_dot(argv[1]);
+  }
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/cpp/task_get_stream.cu

@@ -0,0 +1,65 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+//! \file
+//!
+//! \brief Test the behavior of the get_stream() method of the tasks in the different backends
+
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+__global__ void dummy() {}
+
+void test_stream()
+{
+  // stream context
+  context ctx;
+
+  auto token = ctx.token();
+  EXPECT(token.is_void_interface());
+  auto t = ctx.task(token.write());
+  t.start();
+  cudaStream_t s = t.get_stream();
+  EXPECT(s != nullptr);
+  dummy<<<1, 1, 0, s>>>();
+  t.end();
+  ctx.finalize();
+}
+
+void test_graph()
+{
+  context ctx = graph_ctx();
+
+  auto token = ctx.token();
+  auto t     = ctx.task(token.write());
+  t.start();
+  cudaStream_t s = t.get_stream();
+  // We are not capturing so there is no stream associated
+  EXPECT(s == nullptr);
+  t.end();
+
+  auto t2 = ctx.task(token.rw());
+  t2.enable_capture();
+  t2.start();
+  cudaStream_t s2 = t2.get_stream();
+  // We are capturing so the stream used for capture is associated to the task
+  EXPECT(s2 != nullptr);
+  t2.end();
+
+  ctx.finalize();
+}
+
+int main()
+{
+  test_stream();
+  test_graph();
+  return 0;
+}

cccl_upstream/cudax/test/stf/cpp/test_pick_stream.cu

@@ -0,0 +1,251 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2025 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+//! \file
+//! \brief Test the stream picking functionality using execution place abstraction
+
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+int main()
+{
+  // Get the number of available devices
+  int device_count;
+  cuda_safe_call(cudaGetDeviceCount(&device_count));
+
+  // Create async_resources_handle for stream pool management.
+  // This can be used independently of any CUDASTF context.
+  async_resources_handle resources;
+
+  // Get current device for comparison
+  int current_device;
+  cuda_safe_call(cudaGetDevice(&current_device));
+
+  // ==========================================================================
+  // Test exec_place::pick_stream() - returns cudaStream_t directly
+  // ==========================================================================
+  {
+    exec_place place = exec_place::current_device();
+
+    // pick_stream() returns a cudaStream_t directly (simpler API)
+    cudaStream_t stream = place.pick_stream(resources);
+    EXPECT(stream != nullptr);
+    EXPECT(get_device_from_stream(stream) == current_device);
+
+    // The for_computation parameter is a performance hint (defaults to true).
+    // When true, uses the computation stream pool; when false, uses the
+    // transfer stream pool. Using separate pools can improve overlapping.
+    cudaStream_t compute_stream  = place.pick_stream(resources, true);
+    cudaStream_t transfer_stream = place.pick_stream(resources, false);
+    EXPECT(compute_stream != nullptr);
+    EXPECT(transfer_stream != nullptr);
+  }
+
+  // ==========================================================================
+  // Test exec_place::getStream() - returns augmented_stream with metadata
+  // ==========================================================================
+  {
+    exec_place place = exec_place::current_device();
+
+    // getStream() returns a augmented_stream with additional metadata
+    augmented_stream dstream = place.getStream(resources, true);
+    EXPECT(dstream.stream != nullptr);
+    EXPECT(dstream.dev_id == current_device);
+    EXPECT(get_device_from_stream(dstream.stream) == current_device);
+  }
+
+  // ==========================================================================
+  // Test stream_pool_size() and pick_all_streams()
+  // ==========================================================================
+  {
+    exec_place place = exec_place::current_device();
+
+    // Query the pool size
+    size_t pool_size = place.stream_pool_size(resources);
+    EXPECT(pool_size > 0);
+    EXPECT(pool_size == async_resources_handle::pool_size);
+
+    // Get all streams from the pool as a vector
+    auto all_streams = place.pick_all_streams(resources);
+    EXPECT(all_streams.size() == pool_size);
+
+    // Verify all streams are valid and on the correct device
+    for (cudaStream_t s : all_streams)
+    {
+      EXPECT(s != nullptr);
+      EXPECT(get_device_from_stream(s) == current_device);
+    }
+  }
+
+  // ==========================================================================
+  // Test with multiple devices
+  // ==========================================================================
+  if (device_count > 1)
+  {
+    for (int test_device = 0; test_device < ::std::min(device_count, 2); ++test_device)
+    {
+      exec_place dev_place = exec_place::device(test_device);
+
+      // pick_stream on a specific device
+      cudaStream_t stream = dev_place.pick_stream(resources);
+      EXPECT(stream != nullptr);
+      EXPECT(get_device_from_stream(stream) == test_device);
+
+      // getStream returns more metadata
+      augmented_stream dstream = dev_place.getStream(resources, true);
+      EXPECT(dstream.stream != nullptr);
+      EXPECT(dstream.dev_id == test_device);
+    }
+  }
+
+  // ==========================================================================
+  // Test activate()/deactivate() - generic alternative to cudaSetDevice
+  // These methods can be used without a CUDASTF context
+  // ==========================================================================
+  {
+    // Save initial device
+    int initial_device;
+    cuda_safe_call(cudaGetDevice(&initial_device));
+
+    // Use activate() to switch to current device (no-op but verifies it works)
+    exec_place current_place = exec_place::current_device();
+    {
+      auto active = current_place.activate();
+
+      int after_activate;
+      cuda_safe_call(cudaGetDevice(&after_activate));
+      EXPECT(after_activate == initial_device);
+    }
+    // exec_place_scope destructor restores automatically
+  }
+
+  // Test activate() with multiple devices using RAII
+  if (device_count > 1)
+  {
+    // Save initial device
+    int initial_device;
+    cuda_safe_call(cudaGetDevice(&initial_device));
+
+    // Switch to device 1 using RAII scope
+    {
+      exec_place place1 = exec_place::device(1);
+      auto active       = place1.activate();
+
+      // Verify we're now on device 1
+      int new_device;
+      cuda_safe_call(cudaGetDevice(&new_device));
+      EXPECT(new_device == 1);
+    }
+    // exec_place_scope destructor restores previous device
+
+    // Verify we're back on the initial device
+    int restored_device;
+    cuda_safe_call(cudaGetDevice(&restored_device));
+    EXPECT(restored_device == initial_device);
+
+    // Nested activation test
+    {
+      exec_place place0 = exec_place::device(0);
+      auto active0      = place0.activate();
+
+      int new_device;
+      cuda_safe_call(cudaGetDevice(&new_device));
+      EXPECT(new_device == 0);
+
+      {
+        exec_place place1 = exec_place::device(1);
+        auto active1      = place1.activate();
+
+        cuda_safe_call(cudaGetDevice(&new_device));
+        EXPECT(new_device == 1);
+      }
+      // active1 destroyed, should restore to device 0
+
+      cuda_safe_call(cudaGetDevice(&new_device));
+      EXPECT(new_device == 0);
+    }
+    // active0 destroyed, should restore to initial device
+
+    cuda_safe_call(cudaGetDevice(&restored_device));
+    EXPECT(restored_device == initial_device);
+  }
+
+  // Test that host exec_place activate works (no-op in practice)
+  {
+    exec_place host_place = exec_place::host();
+    auto active           = host_place.activate();
+  }
+
+  // ==========================================================================
+  // Test context stream picking (for comparison)
+  // ==========================================================================
+  {
+    context ctx;
+    // Contexts also have pick_stream() which uses the default execution place
+    cudaStream_t stream = ctx.pick_stream();
+    EXPECT(stream != nullptr);
+    EXPECT(get_device_from_stream(stream) == current_device);
+    ctx.finalize();
+  }
+
+  // ==========================================================================
+  // Test using exec_place::pick_stream with a context's async_resources
+  // When working alongside a context, use ctx.async_resources() to share
+  // the same stream pools between your code and the context's operations.
+  // ==========================================================================
+  {
+    stream_ctx ctx;
+
+    // Get a stream from a specific execution place using the context's resources
+    exec_place place     = exec_place::current_device();
+    cudaStream_t stream1 = place.pick_stream(ctx.async_resources());
+    EXPECT(stream1 != nullptr);
+    EXPECT(get_device_from_stream(stream1) == current_device);
+
+    // This stream comes from the same pool used by ctx internally
+    cudaStream_t stream2 = ctx.pick_stream();
+    EXPECT(stream2 != nullptr);
+
+    // Both methods use the same underlying stream pool
+    // (streams may or may not be the same depending on round-robin selection)
+
+    ctx.finalize();
+  }
+
+  // ==========================================================================
+  // Test with graph context
+  // ==========================================================================
+  {
+    graph_ctx gctx;
+    cudaStream_t stream = gctx.pick_stream();
+    EXPECT(stream != nullptr);
+    EXPECT(get_device_from_stream(stream) == current_device);
+    gctx.finalize();
+  }
+
+  // ==========================================================================
+  // Test context with execution affinity
+  // ==========================================================================
+  if (device_count > 1)
+  {
+    context ctx;
+
+    exec_place dev1_place = exec_place::device(1);
+    ctx.push_affinity(::std::make_shared<exec_place>(dev1_place));
+
+    // Stream should now come from device 1's pool
+    cudaStream_t affinity_stream = ctx.pick_stream();
+    EXPECT(affinity_stream != nullptr);
+    EXPECT(get_device_from_stream(affinity_stream) == 1);
+
+    ctx.finalize();
+  }
+}

cccl_upstream/cudax/test/stf/cpp/test_pick_stream_green_context.cu

@@ -0,0 +1,259 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2025 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+//! \file
+//! \brief Test the pick_stream functionality with green contexts
+
+#include <cuda/experimental/__places/exec/green_context.cuh>
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+// Green contexts are only supported since CUDA 12.4
+#if _CCCL_CTK_AT_LEAST(12, 4)
+
+//! \brief Verify that a stream belongs to the expected green context
+void verify_stream_green_context(cudaStream_t stream, CUgreenCtx expected_g_ctx)
+{
+  // Get the green context associated to that CUDA stream
+  CUgreenCtx stream_cugc;
+  cuda_safe_call(cuStreamGetGreenCtx(CUstream(stream), &stream_cugc));
+  EXPECT(stream_cugc != nullptr);
+
+  CUcontext stream_green_primary;
+  CUcontext expected_green_primary;
+
+  unsigned long long stream_ctxId;
+  unsigned long long expected_ctxId;
+
+  // Convert green contexts to primary contexts and get their ID
+  cuda_safe_call(cuCtxFromGreenCtx(&stream_green_primary, stream_cugc));
+  cuda_safe_call(cuCtxGetId(stream_green_primary, &stream_ctxId));
+
+  cuda_safe_call(cuCtxFromGreenCtx(&expected_green_primary, expected_g_ctx));
+  cuda_safe_call(cuCtxGetId(expected_green_primary, &expected_ctxId));
+
+  // Make sure the stream belongs to the same green context as expected
+  EXPECT(stream_ctxId == expected_ctxId);
+}
+
+#endif // _CCCL_CTK_AT_LEAST(12, 4)
+
+int main()
+{
+#if _CCCL_CTK_BELOW(12, 4)
+  // Green contexts are not supported, skip the test
+  return 0;
+#else // ^^^ _CCCL_CTK_BELOW(12, 4) ^^^ / vvv _CCCL_CTK_AT_LEAST(12, 4) vvv
+
+  // Get current device
+  int current_device;
+  cuda_safe_call(cudaGetDevice(&current_device));
+
+  // Create green context helper with 8 SMs per context
+  const int num_sms = 8;
+  green_context_helper gc(num_sms, current_device);
+
+  // Create async_resources_handle for stream pool management.
+  // This can be used independently of any CUDASTF context.
+  async_resources_handle resources;
+
+  // ==========================================================================
+  // Compare regular device vs green context execution places
+  // ==========================================================================
+  exec_place regular_device_place = exec_place::current_device();
+
+  // pick_stream() returns cudaStream_t directly
+  cudaStream_t device_stream = regular_device_place.pick_stream(resources);
+  EXPECT(device_stream != nullptr);
+  EXPECT(get_device_from_stream(device_stream) == current_device);
+
+  // ==========================================================================
+  // Test green context execution places - each has isolated stream pools
+  // ==========================================================================
+  auto cnt = gc.get_count();
+  if (cnt > 0)
+  {
+    // Test first green context view - demonstrates place-specific stream pools
+    auto view0           = gc.get_view(0);
+    exec_place gc_place0 = exec_place::green_ctx(view0);
+
+    // Green context execution place uses its dedicated stream pool (not shared device pool)
+    cudaStream_t gc_stream = gc_place0.pick_stream(resources);
+    EXPECT(gc_stream != nullptr);
+    EXPECT(get_device_from_stream(gc_stream) == current_device);
+
+    // Verify the stream belongs to the correct green context
+    verify_stream_green_context(gc_stream, view0.g_ctx);
+
+    // Test with multiple views - demonstrates isolation between green contexts
+    if (cnt > 1)
+    {
+      auto view1           = gc.get_view(1);
+      exec_place gc_place1 = exec_place::green_ctx(view1);
+
+      cudaStream_t gc_stream1 = gc_place1.pick_stream(resources);
+      EXPECT(gc_stream1 != nullptr);
+      EXPECT(get_device_from_stream(gc_stream1) == current_device);
+
+      // Each green context has its own isolated stream pool
+      verify_stream_green_context(gc_stream1, view1.g_ctx);
+
+      // Streams from different green context places are isolated
+      EXPECT(gc_stream != gc_stream1);
+    }
+
+    // getStream() provides additional metadata if needed
+    augmented_stream dstream = gc_place0.getStream(resources, true);
+    EXPECT(dstream.stream != nullptr);
+    EXPECT(dstream.dev_id == current_device);
+  }
+
+  // ==========================================================================
+  // Test activate()/deactivate() with green contexts
+  // These methods can be used without a CUDASTF context
+  // ==========================================================================
+  if (cnt > 0)
+  {
+    auto view           = gc.get_view(0);
+    exec_place gc_place = exec_place::green_ctx(view);
+
+    // Save the current CUDA context
+    CUcontext initial_ctx;
+    cuda_safe_call(cuCtxGetCurrent(&initial_ctx));
+    unsigned long long initial_ctx_id;
+    cuda_safe_call(cuCtxGetId(initial_ctx, &initial_ctx_id));
+
+    {
+      // Activate the green context using RAII
+      auto active = gc_place.activate();
+
+      // Verify the current context is now the green context
+      CUcontext current_ctx;
+      cuda_safe_call(cuCtxGetCurrent(&current_ctx));
+
+      // The current context should be the green context's driver context
+      CUcontext green_driver_ctx;
+      cuda_safe_call(cuCtxFromGreenCtx(&green_driver_ctx, view.g_ctx));
+
+      unsigned long long current_ctx_id, green_ctx_id;
+      cuda_safe_call(cuCtxGetId(current_ctx, &current_ctx_id));
+      cuda_safe_call(cuCtxGetId(green_driver_ctx, &green_ctx_id));
+      EXPECT(current_ctx_id == green_ctx_id);
+    }
+    // exec_place_scope destructor restores previous context
+
+    // Verify we're back to the initial context
+    CUcontext restored_ctx;
+    cuda_safe_call(cuCtxGetCurrent(&restored_ctx));
+    unsigned long long restored_ctx_id;
+    cuda_safe_call(cuCtxGetId(restored_ctx, &restored_ctx_id));
+    EXPECT(initial_ctx_id == restored_ctx_id);
+  }
+
+  // Test switching between multiple green contexts using nested RAII
+  if (cnt > 1)
+  {
+    auto view0           = gc.get_view(0);
+    auto view1           = gc.get_view(1);
+    exec_place gc_place0 = exec_place::green_ctx(view0);
+    exec_place gc_place1 = exec_place::green_ctx(view1);
+
+    CUcontext green0_ctx, green1_ctx;
+    cuda_safe_call(cuCtxFromGreenCtx(&green0_ctx, view0.g_ctx));
+    cuda_safe_call(cuCtxFromGreenCtx(&green1_ctx, view1.g_ctx));
+    unsigned long long green0_id, green1_id;
+    cuda_safe_call(cuCtxGetId(green0_ctx, &green0_id));
+    cuda_safe_call(cuCtxGetId(green1_ctx, &green1_id));
+
+    {
+      // Activate first green context
+      auto active0 = gc_place0.activate();
+
+      // Verify we're in green context 0
+      CUcontext current_ctx;
+      cuda_safe_call(cuCtxGetCurrent(&current_ctx));
+      unsigned long long current_id;
+      cuda_safe_call(cuCtxGetId(current_ctx, &current_id));
+      EXPECT(current_id == green0_id);
+
+      {
+        // Switch to second green context (nested)
+        auto active1 = gc_place1.activate();
+
+        // Verify we're now in green context 1
+        cuda_safe_call(cuCtxGetCurrent(&current_ctx));
+        cuda_safe_call(cuCtxGetId(current_ctx, &current_id));
+        EXPECT(current_id == green1_id);
+      }
+      // active1 destroyed, should restore to green context 0
+
+      cuda_safe_call(cuCtxGetCurrent(&current_ctx));
+      cuda_safe_call(cuCtxGetId(current_ctx, &current_id));
+      EXPECT(current_id == green0_id);
+    }
+    // active0 destroyed, restores to original context
+  }
+
+  // ==========================================================================
+  // Test context with green context affinity
+  // ==========================================================================
+  {
+    stream_ctx ctx;
+
+    if (cnt > 0)
+    {
+      // Set affinity to green context execution place
+      auto view           = gc.get_view(0);
+      exec_place gc_place = exec_place::green_ctx(view);
+
+      ctx.push_affinity(::std::make_shared<exec_place>(gc_place));
+
+      // Context pick_stream() respects the green context affinity
+      cudaStream_t stream = ctx.pick_stream();
+      EXPECT(stream != nullptr);
+      EXPECT(get_device_from_stream(stream) == current_device);
+
+      // Verify stream belongs to the green context we set as affinity
+      verify_stream_green_context(stream, view.g_ctx);
+    }
+
+    ctx.finalize();
+  }
+
+  // ==========================================================================
+  // Test graph context with green context affinity
+  // ==========================================================================
+  {
+    graph_ctx gctx;
+
+    if (cnt > 0)
+    {
+      // Set green context affinity for graph context
+      auto view           = gc.get_view(0);
+      exec_place gc_place = exec_place::green_ctx(view);
+
+      gctx.push_affinity(::std::make_shared<exec_place>(gc_place));
+
+      // Graph context also respects the execution place abstraction
+      cudaStream_t graph_stream = gctx.pick_stream();
+      EXPECT(graph_stream != nullptr);
+      EXPECT(get_device_from_stream(graph_stream) == current_device);
+
+      // Verify graph submission stream also respects green context affinity
+      verify_stream_green_context(graph_stream, view.g_ctx);
+    }
+
+    gctx.finalize();
+  }
+
+  return 0;
+#endif // ^^^ _CCCL_CTK_AT_LEAST(12, 4) ^^^
+}

cccl_upstream/cudax/test/stf/cpp/user_streams.cu

@@ -0,0 +1,99 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+/*
+ * In this example, the user provides streams in which the STF model inserts the proper dependencies
+ */
+
+static __global__ void cuda_sleep_kernel(long long int clock_cnt)
+{
+  long long int start_clock  = clock64();
+  long long int clock_offset = 0;
+  while (clock_offset < clock_cnt)
+  {
+    clock_offset = clock64() - start_clock;
+  }
+}
+
+void cuda_sleep(double ms, cudaStream_t stream)
+{
+  int device;
+  cudaGetDevice(&device);
+
+  // cudaDevAttrClockRate: Peak clock frequency in kilohertz;
+  int clock_rate;
+  cudaDeviceGetAttribute(&clock_rate, cudaDevAttrClockRate, device);
+
+  long long int clock_cnt = (long long int) (ms * clock_rate);
+  cuda_sleep_kernel<<<1, 1, 0, stream>>>(clock_cnt);
+}
+
+int main()
+{
+  stream_ctx ctx;
+  double vA, vB, vC, vD;
+  auto A = ctx.logical_data(make_slice(&vA, 1));
+  auto B = ctx.logical_data(make_slice(&vB, 1));
+  auto C = ctx.logical_data(make_slice(&vC, 1));
+  auto D = ctx.logical_data(make_slice(&vD, 1));
+
+  // We are going to submit kernels with the following data accesses, where
+  // K2 and K3 can be executed concurrently, after K1 and been executed, and
+  // before K4 is executed.
+  // K1(Aw); K2(Ar,Bw); K3(Ar, Cw); K4(Br,Cr,Dw);
+
+  // User-provided streams
+  cudaStream_t K1_stream;
+  cudaStream_t K2_stream;
+  cudaStream_t K3_stream;
+  cudaStream_t K4_stream;
+  cudaStreamCreate(&K1_stream);
+  cudaStreamCreate(&K2_stream);
+  cudaStreamCreate(&K3_stream);
+  cudaStreamCreate(&K4_stream);
+
+  // Kernel 1 : A(write)
+  auto k1 = ctx.task(A.rw());
+  k1.set_stream(K1_stream);
+  k1.set_symbol("K1");
+  k1.start();
+  cuda_sleep(500, K1_stream);
+  k1.end();
+
+  // Kernel 2 : A(read) B(write)
+  auto k2 = ctx.task(A.read(), B.write());
+  k2.set_stream(K2_stream);
+  k2.set_symbol("K2");
+  k2.start();
+  cuda_sleep(500, K2_stream);
+  k2.end();
+
+  // Kernel 3 : A(read) C(write)
+  auto k3 = ctx.task(A.read(), C.write());
+  k3.set_stream(K3_stream);
+  k3.set_symbol("K3");
+  k3.start();
+  cuda_sleep(500, K3_stream);
+  k3.end();
+
+  // Kernel 4 : B(read) C(read) D(write)
+  auto k4 = ctx.task(B.read(), C.read(), D.write());
+  k4.set_stream(K4_stream);
+  k4.set_symbol("K4");
+  k4.start();
+  cuda_sleep(500, K4_stream);
+  k4.end();
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/cpp/wait_token.cu

@@ -0,0 +1,60 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+//! \file
+//!
+//! \brief Test ctx.wait() on a token: a blocking, value-less synchronization
+
+#include <cuda/experimental/stf.cuh>
+
+#include <type_traits>
+
+using namespace cuda::experimental::stf;
+
+__global__ void set_value(int* p, int v)
+{
+  *p = v;
+}
+
+template <typename context_t>
+void run()
+{
+  context_t ctx;
+
+  // Externally owned buffer: STF only schedules around it, it never owns it.
+  int* d_val = nullptr;
+  cuda_safe_call(cudaMalloc(&d_val, sizeof(int)));
+
+  auto tok = ctx.token();
+
+  ctx.task(tok.write())->*[=](cudaStream_t s) {
+    set_value<<<1, 1, 0, s>>>(d_val, 42);
+  };
+
+  // wait(token) has no value to materialize: it must return void and only
+  // block the host until the token's producing work has completed.
+  static_assert(::std::is_void_v<decltype(ctx.wait(tok))>, "wait(token) must return void");
+  ctx.wait(tok);
+
+  int h_val = 0;
+  cuda_safe_call(cudaMemcpy(&h_val, d_val, sizeof(int), cudaMemcpyDeviceToHost));
+  _CCCL_ASSERT(h_val == 42, "wait(token) did not synchronize the producing task");
+
+  ctx.finalize();
+  cuda_safe_call(cudaFree(d_val));
+}
+
+int main()
+{
+  run<stream_ctx>();
+  run<graph_ctx>();
+  run<context>();
+  run<stackable_ctx>();
+}

cccl_upstream/cudax/test/stf/cuda-samples/0_Introduction/vectorAdd/README.md

@@ -0,0 +1,69 @@
+# vectorAdd - Vector Addition
+
+## Description
+
+This CUDA Runtime API sample is a very basic sample that implements element by element vector addition. It is the same as the sample illustrating Chapter 3 of the programming guide with some additions like error checking.
+
+## Key Concepts
+
+CUDA Runtime API, Vector Addition
+
+## Supported SM Architectures
+
+[SM 3.5 ](https://developer.nvidia.com/cuda-gpus)  [SM 3.7 ](https://developer.nvidia.com/cuda-gpus)  [SM 5.0 ](https://developer.nvidia.com/cuda-gpus)  [SM 5.2 ](https://developer.nvidia.com/cuda-gpus)  [SM 5.3 ](https://developer.nvidia.com/cuda-gpus)  [SM 6.0 ](https://developer.nvidia.com/cuda-gpus)  [SM 6.1 ](https://developer.nvidia.com/cuda-gpus)  [SM 7.0 ](https://developer.nvidia.com/cuda-gpus)  [SM 7.2 ](https://developer.nvidia.com/cuda-gpus)  [SM 7.5 ](https://developer.nvidia.com/cuda-gpus)  [SM 8.0 ](https://developer.nvidia.com/cuda-gpus)  [SM 8.6 ](https://developer.nvidia.com/cuda-gpus)  [SM 8.7 ](https://developer.nvidia.com/cuda-gpus)
+
+## Supported OSes
+
+Linux, Windows
+
+## Supported CPU Architecture
+
+x86_64, ppc64le, armv7l
+
+## CUDA APIs involved
+
+### [CUDA Runtime API](http://docs.nvidia.com/cuda/cuda-runtime-api/index.html)
+cudaFree, cudaMalloc, cudaGetLastError, cudaMemcpy, cudaGetErrorString
+
+## Prerequisites
+
+Download and install the [CUDA Toolkit 11.6](https://developer.nvidia.com/cuda-downloads) for your corresponding platform.
+
+## Build and Run
+
+### Windows
+The Windows samples are built using the Visual Studio IDE. Solution files (.sln) are provided for each supported version of Visual Studio, using the format:
+```
+*_vs<version>.sln - for Visual Studio <version>
+```
+Each individual sample has its own set of solution files in its directory:
+
+To build/examine all the samples at once, the complete solution files should be used. To build/examine a single sample, the individual sample solution files should be used.
+> **Note:** Some samples require that the Microsoft DirectX SDK (June 2010 or newer) be installed and that the VC++ directory paths are properly set up (**Tools > Options...**). Check DirectX Dependencies section for details."
+
+### Linux
+The Linux samples are built using makefiles. To use the makefiles, change the current directory to the sample directory you wish to build, and run make:
+```
+$ cd <sample_dir>
+$ make
+```
+The samples makefiles can take advantage of certain options:
+*  **TARGET_ARCH=<arch>** - cross-compile targeting a specific architecture. Allowed architectures are x86_64, ppc64le, armv7l.
+    By default, TARGET_ARCH is set to HOST_ARCH. On a x86_64 machine, not setting TARGET_ARCH is the equivalent of setting TARGET_ARCH=x86_64.<br/>
+`$ make TARGET_ARCH=x86_64` <br/> `$ make TARGET_ARCH=ppc64le` <br/> `$ make TARGET_ARCH=armv7l` <br/>
+    See [here](http://docs.nvidia.com/cuda/cuda-samples/index.html#cross-samples) for more details.
+*   **dbg=1** - build with debug symbols
+    ```
+    $ make dbg=1
+    ```
+*   **SMS="A B ..."** - override the SM architectures for which the sample will be built, where `"A B ..."` is a space-delimited list of SM architectures. For example, to generate SASS for SM 50 and SM 60, use `SMS="50 60"`.
+    ```
+    $ make SMS="50 60"
+    ```
+
+*  **HOST_COMPILER=<host_compiler>** - override the default g++ host compiler. See the [Linux Installation Guide](http://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html#system-requirements) for a list of supported host compilers.
+```
+    $ make HOST_COMPILER=g++
+```
+
+## References (for more details)

cccl_upstream/cudax/test/stf/cuda-samples/0_Introduction/vectorAdd/vectorAdd.cu

@@ -0,0 +1,44 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ *  * Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ *  * Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in the
+ *    documentation and/or other materials provided with the distribution.
+ *  * Neither the name of NVIDIA CORPORATION nor the names of its
+ *    contributors may be used to endorse or promote products derived
+ *    from this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+ * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+ * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+ * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+ * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+ * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+ * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+ * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+/**
+ * Vector addition: C = A + B.
+ *
+ * This sample is a very basic sample that implements element by element
+ * vector addition. It is the same as the sample illustrating Chapter 2
+ * of the programming guide with some additions like error checking.
+ */

cccl_upstream/cudax/test/stf/cuda-samples/0_Introduction/vectorAdd/vectorAdd_cudastf.cu

@@ -0,0 +1,149 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ *  * Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ *  * Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in the
+ *    documentation and/or other materials provided with the distribution.
+ *  * Neither the name of NVIDIA CORPORATION nor the names of its
+ *    contributors may be used to endorse or promote products derived
+ *    from this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+ * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+ * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+ * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+ * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+ * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+ * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+ * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+/**
+ * Vector addition: C = A + B.
+ *
+ * This sample is a very basic sample that implements element by element
+ * vector addition. It is the same as the sample illustrating Chapter 2
+ * of the programming guide with some additions like error checking.
+ */
+
+#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+/**
+ * CUDA Kernel Device code
+ *
+ * Computes the vector addition of A and B into C. The 3 vectors have the same
+ * number of elements numElements.
+ */
+__global__ void vectorAdd(const float* A, const float* B, float* C, int numElements)
+{
+  int i = blockDim.x * blockIdx.x + threadIdx.x;
+
+  if (i < numElements)
+  {
+    C[i] = A[i] + B[i] + 0.0f;
+  }
+}
+
+template <typename Ctx>
+void run()
+{
+  Ctx ctx;
+  // Error code to check return values for CUDA calls
+  cudaError_t err = cudaSuccess;
+
+  // Print the vector length to be used, and compute its size
+  int numElements = 50000;
+  size_t size     = numElements * sizeof(float);
+  // printf("[Vector addition of %d elements]\n", numElements);
+
+  // Allocate the host input vector A
+  float* h_A = (float*) malloc(size);
+
+  // Allocate the host input vector B
+  float* h_B = (float*) malloc(size);
+
+  // Allocate the host output vector C
+  float* h_C = (float*) malloc(size);
+
+  // Verify that allocations succeeded
+  if (h_A == NULL || h_B == NULL || h_C == NULL)
+  {
+    fprintf(stderr, "Failed to allocate host vectors!\n");
+    exit(EXIT_FAILURE);
+  }
+
+  // Initialize the host input vectors
+  for (int i = 0; i < numElements; ++i)
+  {
+    h_A[i] = rand() / (float) RAND_MAX;
+    h_B[i] = rand() / (float) RAND_MAX;
+  }
+
+  auto A_handle = ctx.logical_data(h_A, numElements);
+  auto B_handle = ctx.logical_data(h_B, numElements);
+  auto C_handle = ctx.logical_data(h_C, numElements);
+
+  ctx.task(A_handle.read(), B_handle.read(), C_handle.write())->*[&](cudaStream_t stream, auto d_A, auto d_B, auto d_C) {
+    // Launch the Vector Add CUDA Kernel
+    int threadsPerBlock = 256;
+    int blocksPerGrid   = (numElements + threadsPerBlock - 1) / threadsPerBlock;
+    // printf("CUDA kernel launch with %d blocks of %d threads\n", blocksPerGrid, threadsPerBlock);
+    vectorAdd<<<blocksPerGrid, threadsPerBlock, 0, stream>>>(
+      d_A.data_handle(), d_B.data_handle(), d_C.data_handle(), numElements);
+    err = cudaGetLastError();
+
+    if (err != cudaSuccess)
+    {
+      fprintf(stderr, "Failed to launch vectorAdd kernel (error code %s)!\n", cudaGetErrorString(err));
+      exit(EXIT_FAILURE);
+    }
+  };
+
+  ctx.host_launch(A_handle.read(), B_handle.read(), C_handle.read())->*[](auto hA, auto hB, auto hC) {
+    // Verify that the result vector is correct
+    for (int i = 0; i < hC.extent(0); ++i)
+    {
+      if (::std::abs(hA[i] + hB[i] - hC[i]) > 1e-5)
+      {
+        fprintf(stderr, "Result verification failed at element %d!\n", i);
+        exit(EXIT_FAILURE);
+      }
+    }
+  };
+
+  ctx.finalize();
+
+  // Free host memory
+  free(h_A);
+  free(h_B);
+  free(h_C);
+}
+
+/**
+ * Host main routine
+ */
+int main(void)
+{
+  run<stream_ctx>();
+  run<graph_ctx>();
+}

cccl_upstream/cudax/test/stf/cuda-samples/3_CUDA_Features/graphConditionalNodes/graphConditionalNodes.cu

@@ -0,0 +1,367 @@
+/* Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ *  * Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ *  * Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in the
+ *    documentation and/or other materials provided with the distribution.
+ *  * Neither the name of NVIDIA CORPORATION nor the names of its
+ *    contributors may be used to endorse or promote products derived
+ *    from this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+ * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+ * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+ * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+ * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+ * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+ * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+ * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+/*
+ * This file demonstrates the usage of conditional graph nodes with
+ * a series of *simple* example graphs.
+ *
+ * For more information on conditional nodes, see the programming guide:
+ *
+ *   https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#conditional-graph-nodes
+ *
+ */
+
+#include <cuda/experimental/stf.cuh>
+
+#include <cassert>
+#include <cstdio>
+
+#include <cuda_runtime.h>
+
+using namespace cuda::experimental::stf;
+
+#if _CCCL_CTK_AT_LEAST(12, 4)
+
+// This kernel will only be executed if the condition is true
+__global__ void doWhileEmptyKernel(void)
+{
+  printf("GPU: doWhileEmptyKernel()\n");
+  return;
+}
+
+__global__ void doWhileLoopKernel(char* dPtr, cudaGraphConditionalHandle handle)
+{
+  if (--(*dPtr) == 0)
+  {
+    cudaGraphSetConditional(handle, 0);
+  }
+  printf("GPU: counter = %d\n", *dPtr);
+}
+
+void simpleDoWhileGraph(void)
+{
+  cudaGraph_t graph;
+  cudaGraphExec_t graphExec;
+  cudaGraphNode_t conditionalNode;
+
+  // Allocate a byte of device memory to use as input
+  char* dPtr;
+  cuda_safe_call(cudaMalloc((void**) &dPtr, 1));
+
+  printf("simpleDoWhileGraph: Building graph...\n");
+  cuda_safe_call(cudaGraphCreate(&graph, 0));
+
+  cudaGraphConditionalHandle handle;
+  cuda_safe_call(cudaGraphConditionalHandleCreate(&handle, graph, 1, cudaGraphCondAssignDefault));
+
+  cudaGraphNodeParams cParams{};
+  cParams.type               = cudaGraphNodeTypeConditional;
+  cParams.conditional.handle = handle;
+  cParams.conditional.type   = cudaGraphCondTypeWhile;
+  cParams.conditional.size   = 1;
+#  if _CCCL_CTK_AT_LEAST(13, 0)
+  cuda_safe_call(cudaGraphAddNode(&conditionalNode, graph, NULL, NULL, 0, &cParams));
+#  else
+  cuda_safe_call(cudaGraphAddNode(&conditionalNode, graph, NULL, 0, &cParams));
+#  endif
+
+  cudaGraph_t bodyGraph = cParams.conditional.phGraph_out[0];
+
+  cudaStream_t captureStream;
+  cuda_safe_call(cudaStreamCreate(&captureStream));
+
+  cuda_safe_call(
+    cudaStreamBeginCaptureToGraph(captureStream, bodyGraph, nullptr, nullptr, 0, cudaStreamCaptureModeGlobal));
+  doWhileEmptyKernel<<<1, 1, 0, captureStream>>>();
+  doWhileEmptyKernel<<<1, 1, 0, captureStream>>>();
+  doWhileLoopKernel<<<1, 1, 0, captureStream>>>(dPtr, handle);
+  cuda_safe_call(cudaStreamEndCapture(captureStream, nullptr));
+  cuda_safe_call(cudaStreamDestroy(captureStream));
+
+  cuda_safe_call(cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0));
+
+  // Initialize device memory and launch the graph
+  cuda_safe_call(cudaMemset(dPtr, 10, 1)); // Set dPtr to 10
+  printf("Host: Launching graph with loop counter set to 10\n");
+  cuda_safe_call(cudaGraphLaunch(graphExec, 0));
+  cuda_safe_call(cudaDeviceSynchronize());
+
+  // Cleanup
+  cuda_safe_call(cudaGraphExecDestroy(graphExec));
+  cuda_safe_call(cudaGraphDestroy(graph));
+  cuda_safe_call(cudaFree(dPtr));
+
+  printf("simpleDoWhileGraph: Complete\n\n");
+}
+
+void stf_dowhile()
+{
+  stackable_ctx ctx;
+
+  {
+    auto repeat_guard = ctx.repeat_graph_scope(10);
+
+    ctx.task()->*[](cudaStream_t stream) {
+      doWhileEmptyKernel<<<1, 1, 0, stream>>>();
+      doWhileEmptyKernel<<<1, 1, 0, stream>>>();
+    };
+  }
+
+  ctx.finalize();
+  printf("STF do while complete\n\n");
+}
+
+/*
+ * Create a graph containing a conditional while loop using stream capture.
+ * This demonstrates how to insert a conditional node into a stream which is
+ * being captured. The graph consists of a kernel node, A, followed by a
+ * conditional while node, B, followed by a kernel node, D. The conditional
+ * body is populated by a single kernel node, C:
+ *
+ * A -> B [ C ] -> D
+ *
+ * The same kernel will be used for both nodes A and C. This kernel will test
+ * a device memory location and set the condition when the location is non-zero.
+ * We must run the kernel before the loop as well as inside the loop in order
+ * to behave like a while loop as opposed to a do-while loop. We need to evaluate
+ * the device memory location before the conditional node is evaluated in order
+ * to set the condition variable properly. Because we're using a kernel upstream
+ * of the conditional node, there is no need to use the handle default value to
+ * initialize the conditional value.
+ */
+
+__global__ void capturedWhileKernel(char* dPtr, cudaGraphConditionalHandle handle)
+{
+  printf("GPU: counter = %d\n", *dPtr);
+  if (*dPtr)
+  {
+    (*dPtr)--;
+  }
+  cudaGraphSetConditional(handle, *dPtr);
+}
+
+__global__ void capturedWhileEmptyKernel(void)
+{
+  printf("GPU: capturedWhileEmptyKernel()\n");
+  return;
+}
+
+void capturedWhileGraph(void)
+{
+  cudaGraph_t graph;
+  cudaGraphExec_t graphExec;
+
+  cudaStreamCaptureStatus status;
+  const cudaGraphNode_t* dependencies;
+  size_t numDependencies;
+
+  // Allocate a byte of device memory to use as input
+  char* dPtr;
+  cuda_safe_call(cudaMalloc((void**) &dPtr, 1));
+
+  printf("capturedWhileGraph: Building graph...\n");
+  cudaStream_t captureStream;
+  cuda_safe_call(cudaStreamCreate(&captureStream));
+
+  cuda_safe_call(cudaStreamBeginCapture(captureStream, cudaStreamCaptureModeGlobal));
+
+  // Obtain the handle of the graph
+#  if _CCCL_CTK_AT_LEAST(13, 0)
+  cuda_safe_call(cudaStreamGetCaptureInfo(captureStream, &status, NULL, &graph, &dependencies, NULL, &numDependencies));
+#  else
+  cuda_safe_call(cudaStreamGetCaptureInfo(captureStream, &status, NULL, &graph, &dependencies, &numDependencies));
+#  endif
+
+  // Create the conditional handle
+  cudaGraphConditionalHandle handle;
+  cuda_safe_call(cudaGraphConditionalHandleCreate(&handle, graph));
+
+  // Insert kernel node A
+  capturedWhileKernel<<<1, 1, 0, captureStream>>>(dPtr, handle);
+
+  // Obtain the handle for node A
+#  if _CCCL_CTK_AT_LEAST(13, 0)
+  cuda_safe_call(cudaStreamGetCaptureInfo(captureStream, &status, NULL, &graph, &dependencies, NULL, &numDependencies));
+#  else
+  cuda_safe_call(cudaStreamGetCaptureInfo(captureStream, &status, NULL, &graph, &dependencies, &numDependencies));
+#  endif
+
+  // Insert conditional node B
+  cudaGraphNode_t conditionalNode;
+  cudaGraphNodeParams cParams{};
+  cParams.type               = cudaGraphNodeTypeConditional;
+  cParams.conditional.handle = handle;
+  cParams.conditional.type   = cudaGraphCondTypeWhile;
+  cParams.conditional.size   = 1;
+#  if _CCCL_CTK_AT_LEAST(13, 0)
+  cuda_safe_call(cudaGraphAddNode(&conditionalNode, graph, dependencies, NULL, numDependencies, &cParams));
+#  else
+  cuda_safe_call(cudaGraphAddNode(&conditionalNode, graph, dependencies, numDependencies, &cParams));
+#  endif
+
+  cudaGraph_t bodyGraph = cParams.conditional.phGraph_out[0];
+
+  // Update stream capture dependencies to account for the node we manually added
+#  if _CCCL_CTK_AT_LEAST(13, 0)
+  cuda_safe_call(
+    cudaStreamUpdateCaptureDependencies(captureStream, &conditionalNode, NULL, 1, cudaStreamSetCaptureDependencies));
+#  else
+  cuda_safe_call(
+    cudaStreamUpdateCaptureDependencies(captureStream, &conditionalNode, 1, cudaStreamSetCaptureDependencies));
+#  endif
+
+  // Insert kernel node D
+  capturedWhileEmptyKernel<<<1, 1, 0, captureStream>>>();
+
+  cuda_safe_call(cudaStreamEndCapture(captureStream, &graph));
+  cuda_safe_call(cudaStreamDestroy(captureStream));
+
+  // Populate conditional body graph using stream capture
+  cudaStream_t bodyStream;
+  cuda_safe_call(cudaStreamCreate(&bodyStream));
+
+  cuda_safe_call(
+    cudaStreamBeginCaptureToGraph(bodyStream, bodyGraph, nullptr, nullptr, 0, cudaStreamCaptureModeGlobal));
+
+  // Insert kernel node C
+  capturedWhileKernel<<<1, 1, 0, bodyStream>>>(dPtr, handle);
+  cuda_safe_call(cudaStreamEndCapture(bodyStream, nullptr));
+  cuda_safe_call(cudaStreamDestroy(bodyStream));
+
+  cuda_safe_call(cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0));
+
+  // Initialize device memory and launch the graph
+  // Device memory is zero, so the conditional node will not execute
+  cuda_safe_call(cudaMemset(dPtr, 0, 1)); // Set dPtr to 0
+  printf("Host: Launching graph with loop counter set to 0\n");
+  cuda_safe_call(cudaGraphLaunch(graphExec, 0));
+  cuda_safe_call(cudaDeviceSynchronize());
+
+  // Initialize device memory and launch the graph
+  cuda_safe_call(cudaMemset(dPtr, 10, 1)); // Set dPtr to 10
+  printf("Host: Launching graph with loop counter set to 10\n");
+  cuda_safe_call(cudaGraphLaunch(graphExec, 0));
+  cuda_safe_call(cudaDeviceSynchronize());
+
+  // Cleanup
+  cuda_safe_call(cudaGraphExecDestroy(graphExec));
+  cuda_safe_call(cudaGraphDestroy(graph));
+  cuda_safe_call(cudaFree(dPtr));
+
+  printf("capturedWhileGraph: Complete\n\n");
+}
+
+void stf_dowhile_2()
+{
+  stackable_ctx ctx;
+
+  {
+    // We force everything to be a CUDA graph
+    auto scope = ctx.graph_scope();
+
+    // We use a token to ensure that A, B(C) and D are serialized
+    auto t = ctx.token();
+
+    // A
+    ctx.task(t.rw())->*[](cudaStream_t stream) {
+      doWhileEmptyKernel<<<1, 1, 0, stream>>>();
+    };
+
+    // B
+    {
+      auto repeat_guard = ctx.repeat_graph_scope(10);
+
+      // C
+      ctx.task(t.rw())->*[](cudaStream_t stream) {
+        doWhileEmptyKernel<<<1, 1, 0, stream>>>();
+      };
+    }
+
+    // D
+    ctx.task(t.rw())->*[](cudaStream_t stream) {
+      doWhileEmptyKernel<<<1, 1, 0, stream>>>();
+    };
+  }
+
+  ctx.finalize();
+}
+
+void stf_dowhile_2_cuda_kernel()
+{
+  stackable_ctx ctx;
+
+  {
+    // We force everything to be a CUDA graph
+    auto scope = ctx.graph_scope();
+
+    // We use a token to ensure that A, B(C) and D are serialized
+    auto t = ctx.token();
+
+    // A
+    ctx.cuda_kernel(t.rw())->*[]() {
+      return cuda_kernel_desc{doWhileEmptyKernel, 1, 1, 0};
+    };
+
+    // B
+    {
+      auto repeat_guard = ctx.repeat_graph_scope(10);
+
+      // C
+      ctx.cuda_kernel(t.rw())->*[]() {
+        return cuda_kernel_desc{doWhileEmptyKernel, 1, 1, 0};
+      };
+    }
+
+    // D
+    ctx.cuda_kernel(t.rw())->*[]() {
+      return cuda_kernel_desc{doWhileEmptyKernel, 1, 1, 0};
+    };
+  }
+
+  ctx.finalize();
+}
+
+#endif // _CCCL_CTK_AT_LEAST(12, 4)
+
+int main(int, char**)
+{
+#if _CCCL_CTK_BELOW(12, 4)
+  fprintf(stderr, "Waiving test: conditional nodes are only available since CUDA 12.4.\n");
+  return 0;
+#else
+  simpleDoWhileGraph();
+  stf_dowhile();
+
+  capturedWhileGraph();
+  stf_dowhile_2();
+
+  // same as stf_dowhile_2 but uses cuda_kernel
+  stf_dowhile_2_cuda_kernel();
+
+  return 0;
+#endif // _CCCL_CTK_AT_LEAST(12, 4)
+}

cccl_upstream/cudax/test/stf/cuda-samples/3_CUDA_Features/jacobiCudaGraphs/README.md

@@ -0,0 +1,69 @@
+# jacobiCudaGraphs - Jacobi CUDA Graphs
+
+## Description
+
+Demonstrates Instantiated CUDA Graph Update with Jacobi Iterative Method using cudaGraphExecKernelNodeSetParams() and cudaGraphExecUpdate() approach.
+
+## Key Concepts
+
+CUDA Graphs, Stream Capture, Instantiated CUDA Graph Update, Cooperative Groups
+
+## Supported SM Architectures
+
+[SM 3.5 ](https://developer.nvidia.com/cuda-gpus)  [SM 3.7 ](https://developer.nvidia.com/cuda-gpus)  [SM 5.0 ](https://developer.nvidia.com/cuda-gpus)  [SM 5.2 ](https://developer.nvidia.com/cuda-gpus)  [SM 5.3 ](https://developer.nvidia.com/cuda-gpus)  [SM 6.0 ](https://developer.nvidia.com/cuda-gpus)  [SM 6.1 ](https://developer.nvidia.com/cuda-gpus)  [SM 7.0 ](https://developer.nvidia.com/cuda-gpus)  [SM 7.2 ](https://developer.nvidia.com/cuda-gpus)  [SM 7.5 ](https://developer.nvidia.com/cuda-gpus)  [SM 8.0 ](https://developer.nvidia.com/cuda-gpus)  [SM 8.6 ](https://developer.nvidia.com/cuda-gpus)  [SM 8.7 ](https://developer.nvidia.com/cuda-gpus)
+
+## Supported OSes
+
+Linux, Windows
+
+## Supported CPU Architecture
+
+x86_64, ppc64le, armv7l
+
+## CUDA APIs involved
+
+### [CUDA Runtime API](http://docs.nvidia.com/cuda/cuda-runtime-api/index.html)
+cudaGraphAddMemsetNode, cudaStreamCreateWithFlags, cudaMemcpyAsync, cudaMallocHost, cudaPitchedPtr, cudaGraphCreate, cudaMalloc, cudaPos, cudaGraphAddMemcpyNode, cudaStreamEndCapture, cudaGraphExecDestroy, cudaStreamBeginCapture, cudaGraphExecKernelNodeSetParams, cudaStreamSynchronize, cudaGraphLaunch, cudaFree, cudaGraphInstantiate, cudaExtent, cudaMemsetAsync, cudaFreeHost, cudaGraphAddKernelNode, cudaGraphExecUpdate
+
+## Prerequisites
+
+Download and install the [CUDA Toolkit 11.6](https://developer.nvidia.com/cuda-downloads) for your corresponding platform.
+
+## Build and Run
+
+### Windows
+The Windows samples are built using the Visual Studio IDE. Solution files (.sln) are provided for each supported version of Visual Studio, using the format:
+```
+*_vs<version>.sln - for Visual Studio <version>
+```
+Each individual sample has its own set of solution files in its directory:
+
+To build/examine all the samples at once, the complete solution files should be used. To build/examine a single sample, the individual sample solution files should be used.
+> **Note:** Some samples require that the Microsoft DirectX SDK (June 2010 or newer) be installed and that the VC++ directory paths are properly set up (**Tools > Options...**). Check DirectX Dependencies section for details."
+
+### Linux
+The Linux samples are built using makefiles. To use the makefiles, change the current directory to the sample directory you wish to build, and run make:
+```
+$ cd <sample_dir>
+$ make
+```
+The samples makefiles can take advantage of certain options:
+*  **TARGET_ARCH=<arch>** - cross-compile targeting a specific architecture. Allowed architectures are x86_64, ppc64le, armv7l.
+    By default, TARGET_ARCH is set to HOST_ARCH. On a x86_64 machine, not setting TARGET_ARCH is the equivalent of setting TARGET_ARCH=x86_64.<br/>
+`$ make TARGET_ARCH=x86_64` <br/> `$ make TARGET_ARCH=ppc64le` <br/> `$ make TARGET_ARCH=armv7l` <br/>
+    See [here](http://docs.nvidia.com/cuda/cuda-samples/index.html#cross-samples) for more details.
+*   **dbg=1** - build with debug symbols
+    ```
+    $ make dbg=1
+    ```
+*   **SMS="A B ..."** - override the SM architectures for which the sample will be built, where `"A B ..."` is a space-delimited list of SM architectures. For example, to generate SASS for SM 50 and SM 60, use `SMS="50 60"`.
+    ```
+    $ make SMS="50 60"
+    ```
+
+*  **HOST_COMPILER=<host_compiler>** - override the default g++ host compiler. See the [Linux Installation Guide](http://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html#system-requirements) for a list of supported host compilers.
+```
+    $ make HOST_COMPILER=g++
+```
+
+## References (for more details)

cccl_upstream/cudax/test/stf/cuda-samples/3_CUDA_Features/jacobiCudaGraphs/jacobi.cu

@@ -0,0 +1,638 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ *  * Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ *  * Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in the
+ *    documentation and/or other materials provided with the distribution.
+ *  * Neither the name of NVIDIA CORPORATION nor the names of its
+ *    contributors may be used to endorse or promote products derived
+ *    from this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+ * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+ * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+ * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+ * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+ * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+ * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+ * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+// This sample demonstrates Instantiated CUDA Graph Update
+// with Jacobi Iterative Method in 3 different methods:
+// 1 - JacobiMethodGpuCudaGraphExecKernelSetParams() - CUDA Graph with
+// cudaGraphExecKernelNodeSetParams() 2 - JacobiMethodGpuCudaGraphExecUpdate() -
+// CUDA Graph with cudaGraphExecUpdate() 3 - JacobiMethodGpu() - Non CUDA Graph
+// method
+
+// Jacobi method on a linear system A*x = b,
+// where A is diagonally dominant and the exact solution consists
+// of all ones.
+
+#include <cuda/experimental/__stf/utility/cuda_safe_call.cuh>
+
+#include <cooperative_groups.h>
+
+using cuda::experimental::stf::cuda_safe_call;
+
+#define N_ROWS 512
+
+namespace cg = cooperative_groups;
+
+// 8 Rows of square-matrix A processed by each CTA.
+// This can be max 32 and only power of 2 (i.e., 2/4/8/16/32).
+#define ROWS_PER_CTA 8
+
+#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 600
+#else
+__device__ double atomicAdd(double* address, double val)
+{
+  unsigned long long int* address_as_ull = (unsigned long long int*) address;
+  unsigned long long int old             = *address_as_ull, assumed;
+
+  do
+  {
+    assumed = old;
+    old     = atomicCAS(address_as_ull, assumed, __double_as_longlong(val + __longlong_as_double(assumed)));
+
+    // Note: uses integer comparison to avoid hang in case of NaN (since NaN !=
+    // NaN)
+  } while (assumed != old);
+
+  return __longlong_as_double(old);
+}
+#endif
+
+// creates N_ROWS x N_ROWS matrix A with N_ROWS+1 on the diagonal and 1
+// elsewhere. The elements of the right hand side b all equal 2*n, hence the
+// exact solution x to A*x = b is a vector of ones.
+void createLinearSystem(float* A, double* b)
+{
+  int i, j;
+  for (i = 0; i < N_ROWS; i++)
+  {
+    b[i] = 2.0 * N_ROWS;
+    for (j = 0; j < N_ROWS; j++)
+    {
+      A[i * N_ROWS + j] = 1.0;
+    }
+    A[i * N_ROWS + i] = N_ROWS + 1.0;
+  }
+}
+
+static __global__ void
+JacobiMethod(const float* A, const double* b, const float conv_threshold, double* x, double* x_new, double* sum)
+{
+  // Handle to thread block group
+  cg::thread_block cta = cg::this_thread_block();
+  __shared__ double x_shared[N_ROWS]; // N_ROWS == n
+  __shared__ double b_shared[ROWS_PER_CTA + 1];
+
+  for (int i = threadIdx.x; i < N_ROWS; i += blockDim.x)
+  {
+    x_shared[i] = x[i];
+  }
+
+  if (threadIdx.x < ROWS_PER_CTA)
+  {
+    int k = threadIdx.x;
+    for (int i = k + (blockIdx.x * ROWS_PER_CTA); (k < ROWS_PER_CTA) && (i < N_ROWS);
+         k += ROWS_PER_CTA, i += ROWS_PER_CTA)
+    {
+      b_shared[i % (ROWS_PER_CTA + 1)] = b[i];
+    }
+  }
+
+  cg::sync(cta);
+
+  cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);
+
+  for (int k = 0, i = blockIdx.x * ROWS_PER_CTA; (k < ROWS_PER_CTA) && (i < N_ROWS); k++, i++)
+  {
+    double rowThreadSum = 0.0;
+    for (int j = threadIdx.x; j < N_ROWS; j += blockDim.x)
+    {
+      rowThreadSum += (A[i * N_ROWS + j] * x_shared[j]);
+    }
+
+    for (int offset = tile32.size() / 2; offset > 0; offset /= 2)
+    {
+      rowThreadSum += tile32.shfl_down(rowThreadSum, offset);
+    }
+
+    if (tile32.thread_rank() == 0)
+    {
+      atomicAdd(&b_shared[i % (ROWS_PER_CTA + 1)], -rowThreadSum);
+    }
+  }
+
+  cg::sync(cta);
+
+  if (threadIdx.x < ROWS_PER_CTA)
+  {
+    cg::thread_block_tile<ROWS_PER_CTA> tile8 = cg::tiled_partition<ROWS_PER_CTA>(cta);
+    double temp_sum                           = 0.0;
+
+    int k = threadIdx.x;
+
+    for (int i = k + (blockIdx.x * ROWS_PER_CTA); (k < ROWS_PER_CTA) && (i < N_ROWS);
+         k += ROWS_PER_CTA, i += ROWS_PER_CTA)
+    {
+      double dx = b_shared[i % (ROWS_PER_CTA + 1)];
+      dx /= A[i * N_ROWS + i];
+
+      x_new[i] = (x_shared[i] + dx);
+      temp_sum += fabs(dx);
+    }
+
+    for (int offset = tile8.size() / 2; offset > 0; offset /= 2)
+    {
+      temp_sum += tile8.shfl_down(temp_sum, offset);
+    }
+
+    if (tile8.thread_rank() == 0)
+    {
+      atomicAdd(sum, temp_sum);
+    }
+  }
+}
+
+// Thread block size for finalError kernel should be multiple of 32
+static __global__ void finalError(double* x, double* g_sum)
+{
+  // Handle to thread block group
+  cg::thread_block cta = cg::this_thread_block();
+  extern __shared__ double warpSum[];
+  double sum = 0.0;
+
+  int globalThreadId = blockIdx.x * blockDim.x + threadIdx.x;
+
+  for (int i = globalThreadId; i < N_ROWS; i += blockDim.x * gridDim.x)
+  {
+    double d = x[i] - 1.0;
+    sum += fabs(d);
+  }
+
+  cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);
+
+  for (int offset = tile32.size() / 2; offset > 0; offset /= 2)
+  {
+    sum += tile32.shfl_down(sum, offset);
+  }
+
+  if (tile32.thread_rank() == 0)
+  {
+    warpSum[threadIdx.x / warpSize] = sum;
+  }
+
+  cg::sync(cta);
+
+  double blockSum = 0.0;
+  if (threadIdx.x < (blockDim.x / warpSize))
+  {
+    blockSum = warpSum[threadIdx.x];
+  }
+
+  if (threadIdx.x < 32)
+  {
+    for (int offset = tile32.size() / 2; offset > 0; offset /= 2)
+    {
+      blockSum += tile32.shfl_down(blockSum, offset);
+    }
+    if (tile32.thread_rank() == 0)
+    {
+      atomicAdd(g_sum, blockSum);
+    }
+  }
+}
+
+// Run the Jacobi method for A*x = b on GPU with CUDA Graph -
+// cudaGraphExecKernelNodeSetParams().
+double JacobiMethodGpuCudaGraphExecKernelSetParams(
+  const float* A,
+  const double* b,
+  float conv_threshold,
+  const int max_iter,
+  double* x,
+  double* x_new,
+  cudaStream_t stream)
+{
+  // CTA size
+  dim3 nthreads(256, 1, 1);
+  // grid size
+  dim3 nblocks((N_ROWS / ROWS_PER_CTA) + 2, 1, 1);
+  cudaGraph_t graph;
+  cudaGraphExec_t graphExec = NULL;
+
+  double sum    = 0.0;
+  double* d_sum = NULL;
+  cuda_safe_call(cudaMalloc(&d_sum, sizeof(double)));
+
+  std::vector<cudaGraphNode_t> nodeDependencies;
+  cudaGraphNode_t memcpyNode, jacobiKernelNode, memsetNode;
+  cudaMemcpy3DParms memcpyParams;
+  cudaMemsetParams memsetParams;
+
+  memsetParams.dst   = (void*) d_sum;
+  memsetParams.value = 0;
+  memsetParams.pitch = 0;
+  // elementSize can be max 4 bytes, so we take sizeof(float) and width=2
+  memsetParams.elementSize = sizeof(float);
+  memsetParams.width       = 2;
+  memsetParams.height      = 1;
+
+  cuda_safe_call(cudaGraphCreate(&graph, 0));
+  cuda_safe_call(cudaGraphAddMemsetNode(&memsetNode, graph, NULL, 0, &memsetParams));
+  nodeDependencies.push_back(memsetNode);
+
+  cudaKernelNodeParams NodeParams0, NodeParams1;
+  NodeParams0.func           = (void*) JacobiMethod;
+  NodeParams0.gridDim        = nblocks;
+  NodeParams0.blockDim       = nthreads;
+  NodeParams0.sharedMemBytes = 0;
+  void* kernelArgs0[6]       = {
+    (void*) &A,
+    const_cast<double**>(&b),
+    const_cast<float*>(&conv_threshold),
+    (void*) &x,
+    (void*) &x_new,
+    (void*) &d_sum};
+  NodeParams0.kernelParams = kernelArgs0;
+  NodeParams0.extra        = NULL;
+
+  cuda_safe_call(
+    cudaGraphAddKernelNode(&jacobiKernelNode, graph, nodeDependencies.data(), nodeDependencies.size(), &NodeParams0));
+
+  nodeDependencies.clear();
+  nodeDependencies.push_back(jacobiKernelNode);
+
+  memcpyParams.srcArray = NULL;
+  memcpyParams.srcPos   = make_cudaPos(0, 0, 0);
+  memcpyParams.srcPtr   = make_cudaPitchedPtr(d_sum, sizeof(double), 1, 1);
+  memcpyParams.dstArray = NULL;
+  memcpyParams.dstPos   = make_cudaPos(0, 0, 0);
+  memcpyParams.dstPtr   = make_cudaPitchedPtr(&sum, sizeof(double), 1, 1);
+  memcpyParams.extent   = make_cudaExtent(sizeof(double), 1, 1);
+  memcpyParams.kind     = cudaMemcpyDeviceToHost;
+
+  cuda_safe_call(
+    cudaGraphAddMemcpyNode(&memcpyNode, graph, nodeDependencies.data(), nodeDependencies.size(), &memcpyParams));
+
+  cuda_safe_call(cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0));
+
+  NodeParams1.func           = (void*) JacobiMethod;
+  NodeParams1.gridDim        = nblocks;
+  NodeParams1.blockDim       = nthreads;
+  NodeParams1.sharedMemBytes = 0;
+  void* kernelArgs1[6]       = {
+    (void*) &A,
+    const_cast<double**>(&b),
+    const_cast<float*>(&conv_threshold),
+    (void*) &x_new,
+    (void*) &x,
+    (void*) &d_sum};
+  NodeParams1.kernelParams = kernelArgs1;
+  NodeParams1.extra        = NULL;
+
+  int k = 0;
+  for (k = 0; k < max_iter; k++)
+  {
+    cuda_safe_call(
+      cudaGraphExecKernelNodeSetParams(graphExec, jacobiKernelNode, ((k & 1) == 0) ? &NodeParams0 : &NodeParams1));
+    cuda_safe_call(cudaGraphLaunch(graphExec, stream));
+    cuda_safe_call(cudaStreamSynchronize(stream));
+
+    if (sum <= conv_threshold)
+    {
+      cuda_safe_call(cudaMemsetAsync(d_sum, 0, sizeof(double), stream));
+      nblocks.x            = (N_ROWS / nthreads.x) + 1;
+      size_t sharedMemSize = ((nthreads.x / 32) + 1) * sizeof(double);
+      if ((k & 1) == 0)
+      {
+        finalError<<<nblocks, nthreads, sharedMemSize, stream>>>(x_new, d_sum);
+      }
+      else
+      {
+        finalError<<<nblocks, nthreads, sharedMemSize, stream>>>(x, d_sum);
+      }
+
+      cuda_safe_call(cudaMemcpyAsync(&sum, d_sum, sizeof(double), cudaMemcpyDeviceToHost, stream));
+      cuda_safe_call(cudaStreamSynchronize(stream));
+      // printf("GPU iterations : %d\n", k + 1);
+      // printf("GPU error: %.3e\n", sum);
+      break;
+    }
+  }
+
+  cuda_safe_call(cudaFree(d_sum));
+  return sum;
+}
+
+// Run the Jacobi method for A*x = b on GPU with Instantiated CUDA Graph Update
+// API - cudaGraphExecUpdate().
+double JacobiMethodGpuCudaGraphExecUpdate(
+  const float* A,
+  const double* b,
+  const float conv_threshold,
+  const int max_iter,
+  double* x,
+  double* x_new,
+  cudaStream_t stream)
+{
+  // CTA size
+  dim3 nthreads(256, 1, 1);
+  // grid size
+  dim3 nblocks((N_ROWS / ROWS_PER_CTA) + 2, 1, 1);
+  cudaGraph_t graph;
+  cudaGraphExec_t graphExec = NULL;
+
+  double sum = 0.0;
+  double* d_sum;
+  cuda_safe_call(cudaMalloc(&d_sum, sizeof(double)));
+
+  int k = 0;
+  for (k = 0; k < max_iter; k++)
+  {
+    cuda_safe_call(cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal));
+    cuda_safe_call(cudaMemsetAsync(d_sum, 0, sizeof(double), stream));
+    if ((k & 1) == 0)
+    {
+      JacobiMethod<<<nblocks, nthreads, 0, stream>>>(A, b, conv_threshold, x, x_new, d_sum);
+    }
+    else
+    {
+      JacobiMethod<<<nblocks, nthreads, 0, stream>>>(A, b, conv_threshold, x_new, x, d_sum);
+    }
+    cuda_safe_call(cudaMemcpyAsync(&sum, d_sum, sizeof(double), cudaMemcpyDeviceToHost, stream));
+    cuda_safe_call(cudaStreamEndCapture(stream, &graph));
+
+    if (graphExec == NULL)
+    {
+      cuda_safe_call(cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0));
+    }
+    else
+    {
+      cudaGraphExecUpdateResult updateResult_out;
+      cuda_safe_call(cudaGraphExecUpdate(graphExec, graph, NULL, &updateResult_out));
+      if (updateResult_out != cudaGraphExecUpdateSuccess)
+      {
+        if (graphExec != NULL)
+        {
+          cuda_safe_call(cudaGraphExecDestroy(graphExec));
+        }
+        printf("k = %d graph update failed with error - %d\n", k, updateResult_out);
+        cuda_safe_call(cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0));
+      }
+    }
+    cuda_safe_call(cudaGraphLaunch(graphExec, stream));
+    cuda_safe_call(cudaStreamSynchronize(stream));
+
+    if (sum <= conv_threshold)
+    {
+      cuda_safe_call(cudaMemsetAsync(d_sum, 0, sizeof(double), stream));
+      nblocks.x            = (N_ROWS / nthreads.x) + 1;
+      size_t sharedMemSize = ((nthreads.x / 32) + 1) * sizeof(double);
+      if ((k & 1) == 0)
+      {
+        finalError<<<nblocks, nthreads, sharedMemSize, stream>>>(x_new, d_sum);
+      }
+      else
+      {
+        finalError<<<nblocks, nthreads, sharedMemSize, stream>>>(x, d_sum);
+      }
+
+      cuda_safe_call(cudaMemcpyAsync(&sum, d_sum, sizeof(double), cudaMemcpyDeviceToHost, stream));
+      cuda_safe_call(cudaStreamSynchronize(stream));
+      // printf("GPU iterations : %d\n", k + 1);
+      // printf("GPU error: %.3e\n", sum);
+      break;
+    }
+  }
+
+  cuda_safe_call(cudaFree(d_sum));
+  return sum;
+}
+
+// Run the Jacobi method for A*x = b on GPU without CUDA Graph.
+double JacobiMethodGpu(
+  const float* A,
+  const double* b,
+  const float conv_threshold,
+  const int max_iter,
+  double* x,
+  double* x_new,
+  cudaStream_t stream)
+{
+  // CTA size
+  dim3 nthreads(256, 1, 1);
+  // grid size
+  dim3 nblocks((N_ROWS / ROWS_PER_CTA) + 2, 1, 1);
+
+  double sum = 0.0;
+  double* d_sum;
+  cuda_safe_call(cudaMalloc(&d_sum, sizeof(double)));
+  int k = 0;
+
+  for (k = 0; k < max_iter; k++)
+  {
+    cuda_safe_call(cudaMemsetAsync(d_sum, 0, sizeof(double), stream));
+    if ((k & 1) == 0)
+    {
+      JacobiMethod<<<nblocks, nthreads, 0, stream>>>(A, b, conv_threshold, x, x_new, d_sum);
+    }
+    else
+    {
+      JacobiMethod<<<nblocks, nthreads, 0, stream>>>(A, b, conv_threshold, x_new, x, d_sum);
+    }
+    cuda_safe_call(cudaMemcpyAsync(&sum, d_sum, sizeof(double), cudaMemcpyDeviceToHost, stream));
+    cuda_safe_call(cudaStreamSynchronize(stream));
+
+    if (sum <= conv_threshold)
+    {
+      cuda_safe_call(cudaMemsetAsync(d_sum, 0, sizeof(double), stream));
+      nblocks.x            = (N_ROWS / nthreads.x) + 1;
+      size_t sharedMemSize = ((nthreads.x / 32) + 1) * sizeof(double);
+      if ((k & 1) == 0)
+      {
+        finalError<<<nblocks, nthreads, sharedMemSize, stream>>>(x_new, d_sum);
+      }
+      else
+      {
+        finalError<<<nblocks, nthreads, sharedMemSize, stream>>>(x, d_sum);
+      }
+
+      cuda_safe_call(cudaMemcpyAsync(&sum, d_sum, sizeof(double), cudaMemcpyDeviceToHost, stream));
+      cuda_safe_call(cudaStreamSynchronize(stream));
+      // printf("GPU iterations : %d\n", k + 1);
+      // printf("GPU error: %.3e\n", sum);
+      break;
+    }
+  }
+
+  cuda_safe_call(cudaFree(d_sum));
+  return sum;
+}
+
+// Run the Jacobi method for A*x = b on CPU.
+void JacobiMethodCPU(float* A, double* b, float conv_threshold, int max_iter, int* num_iter, double* x)
+{
+  double* x_new;
+  x_new = (double*) calloc(N_ROWS, sizeof(double));
+  int k;
+
+  for (k = 0; k < max_iter; k++)
+  {
+    double sum = 0.0;
+    for (int i = 0; i < N_ROWS; i++)
+    {
+      double temp_dx = b[i];
+      for (int j = 0; j < N_ROWS; j++)
+      {
+        temp_dx -= A[i * N_ROWS + j] * x[j];
+      }
+      temp_dx /= A[i * N_ROWS + i];
+      x_new[i] += temp_dx;
+      sum += fabs(temp_dx);
+    }
+
+    for (int i = 0; i < N_ROWS; i++)
+    {
+      x[i] = x_new[i];
+    }
+
+    if (sum <= conv_threshold)
+    {
+      break;
+    }
+  }
+  *num_iter = k + 1;
+  free(x_new);
+}
+
+int main()
+{
+  //  if (checkCmdLineFlag(argc, (const char **)argv, "help")) {
+  //    printf("Command line: jacobiCudaGraphs [-option]\n");
+  //    printf("Valid options:\n");
+  //    printf(
+  //        "-gpumethod=<0,1 or 2>  : 0 - [Default] "
+  //        "JacobiMethodGpuCudaGraphExecKernelSetParams\n");
+  //    printf("                       : 1 - JacobiMethodGpuCudaGraphExecUpdate\n");
+  //    printf("                       : 2 - JacobiMethodGpu - Non CUDA Graph\n");
+  //    printf("-device=device_num     : cuda device id");
+  //    printf("-help         : Output a help message\n");
+  //    exit(EXIT_SUCCESS);
+  //  }
+  //
+  int gpumethod = 0;
+  //  if (checkCmdLineFlag(argc, (const char **)argv, "gpumethod")) {
+  //    gpumethod = getCmdLineArgumentInt(argc, (const char **)argv, "gpumethod");
+  //
+  //    if (gpumethod < 0 || gpumethod > 2) {
+  //      printf("Error: gpumethod must be 0 or 1 or 2, gpumethod=%d is invalid\n",
+  //             gpumethod);
+  //      exit(EXIT_SUCCESS);
+  //    }
+  //  }
+
+  //  int dev = findCudaDevice(argc, (const char **)argv);
+  // int dev = 0;
+
+  double* b = NULL;
+  float* A  = NULL;
+  cuda_safe_call(cudaMallocHost(&b, N_ROWS * sizeof(double)));
+  memset(b, 0, N_ROWS * sizeof(double));
+  cuda_safe_call(cudaMallocHost(&A, N_ROWS * N_ROWS * sizeof(float)));
+  memset(A, 0, N_ROWS * N_ROWS * sizeof(float));
+
+  createLinearSystem(A, b);
+  double* x = NULL;
+  // start with array of all zeroes
+  x = (double*) calloc(N_ROWS, sizeof(double));
+
+  float conv_threshold = 1.0e-2;
+  int max_iter         = 4 * N_ROWS * N_ROWS;
+  int cnt              = 0;
+
+  //  // create timer
+  //  StopWatchInterface *timerCPU = NULL, *timerGpu = NULL;
+  //  sdkCreateTimer(&timerCPU);
+  //
+  //  sdkStartTimer(&timerCPU);
+  JacobiMethodCPU(A, b, conv_threshold, max_iter, &cnt, x);
+
+  double sum = 0.0;
+  // Compute error
+  for (int i = 0; i < N_ROWS; i++)
+  {
+    double d = x[i] - 1.0;
+    sum += fabs(d);
+  }
+  //  sdkStopTimer(&timerCPU);
+  // printf("CPU iterations : %d\n", cnt);
+  // printf("CPU error: %.3e\n", sum);
+  //  printf("CPU Processing time: %f (ms)\n", sdkGetTimerValue(&timerCPU));
+
+  float* d_A;
+  double *d_b, *d_x, *d_x_new;
+  cudaStream_t stream1;
+  cuda_safe_call(cudaStreamCreateWithFlags(&stream1, cudaStreamNonBlocking));
+  cuda_safe_call(cudaMalloc(&d_b, sizeof(double) * N_ROWS));
+  cuda_safe_call(cudaMalloc(&d_A, sizeof(float) * N_ROWS * N_ROWS));
+  cuda_safe_call(cudaMalloc(&d_x, sizeof(double) * N_ROWS));
+  cuda_safe_call(cudaMalloc(&d_x_new, sizeof(double) * N_ROWS));
+
+  cuda_safe_call(cudaMemsetAsync(d_x, 0, sizeof(double) * N_ROWS, stream1));
+  cuda_safe_call(cudaMemsetAsync(d_x_new, 0, sizeof(double) * N_ROWS, stream1));
+  cuda_safe_call(cudaMemcpyAsync(d_A, A, sizeof(float) * N_ROWS * N_ROWS, cudaMemcpyHostToDevice, stream1));
+  cuda_safe_call(cudaMemcpyAsync(d_b, b, sizeof(double) * N_ROWS, cudaMemcpyHostToDevice, stream1));
+
+  //  sdkCreateTimer(&timerGpu);
+  //  sdkStartTimer(&timerGpu);
+
+  double sumGPU = 0.0;
+  if (gpumethod == 0)
+  {
+    sumGPU = JacobiMethodGpuCudaGraphExecKernelSetParams(d_A, d_b, conv_threshold, max_iter, d_x, d_x_new, stream1);
+  }
+  else if (gpumethod == 1)
+  {
+    sumGPU = JacobiMethodGpuCudaGraphExecUpdate(d_A, d_b, conv_threshold, max_iter, d_x, d_x_new, stream1);
+  }
+  else if (gpumethod == 2)
+  {
+    sumGPU = JacobiMethodGpu(d_A, d_b, conv_threshold, max_iter, d_x, d_x_new, stream1);
+  }
+
+  //  sdkStopTimer(&timerGpu);
+  //  printf("GPU Processing time: %f (ms)\n", sdkGetTimerValue(&timerGpu));
+
+  cuda_safe_call(cudaFree(d_b));
+  cuda_safe_call(cudaFree(d_A));
+  cuda_safe_call(cudaFree(d_x));
+  cuda_safe_call(cudaFree(d_x_new));
+
+  cuda_safe_call(cudaFreeHost(A));
+  cuda_safe_call(cudaFreeHost(b));
+
+  // printf("&&&& jacobiCudaGraphs %s\n", (fabs(sum - sumGPU) < conv_threshold) ? "PASSED" : "FAILED");
+
+  return (fabs(sum - sumGPU) < conv_threshold) ? EXIT_SUCCESS : EXIT_FAILURE;
+}

cccl_upstream/cudax/test/stf/cuda-samples/3_CUDA_Features/jacobiCudaGraphs/jacobi_cudastf.cu

@@ -0,0 +1,376 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ *  * Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ *  * Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in the
+ *    documentation and/or other materials provided with the distribution.
+ *  * Neither the name of NVIDIA CORPORATION nor the names of its
+ *    contributors may be used to endorse or promote products derived
+ *    from this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+ * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+ * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+ * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+ * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+ * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+ * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+ * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+// Jacobi method on a linear system A*x = b,
+// where A is diagonally dominant and the exact solution consists
+// of all ones.
+
+#include <cuda/experimental/__stf/internal/scalar_interface.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+#define N_ROWS 512
+
+namespace cg = cooperative_groups;
+
+using namespace cuda::experimental::stf;
+
+// 8 Rows of square-matrix A processed by each CTA.
+// This can be max 32 and only power of 2 (i.e., 2/4/8/16/32).
+#define ROWS_PER_CTA 8
+
+// creates N_ROWS x N_ROWS matrix A with N_ROWS+1 on the diagonal and 1
+// elsewhere. The elements of the right hand side b all equal 2*n, hence the
+// exact solution x to A*x = b is a vector of ones.
+void createLinearSystem(float* A, double* b)
+{
+  int i, j;
+  for (i = 0; i < N_ROWS; i++)
+  {
+    b[i] = 2.0 * N_ROWS;
+    for (j = 0; j < N_ROWS; j++)
+    {
+      A[i * N_ROWS + j] = 1.0;
+    }
+    A[i * N_ROWS + i] = N_ROWS + 1.0;
+  }
+}
+
+static __global__ void
+JacobiMethod(const float* A, const double* b, const float conv_threshold, double* x, double* x_new, double* sum)
+{
+  // Handle to thread block group
+  cg::thread_block cta = cg::this_thread_block();
+  __shared__ double x_shared[N_ROWS]; // N_ROWS == n
+  __shared__ double b_shared[ROWS_PER_CTA + 1];
+
+  for (int i = threadIdx.x; i < N_ROWS; i += blockDim.x)
+  {
+    x_shared[i] = x[i];
+  }
+
+  if (threadIdx.x < ROWS_PER_CTA)
+  {
+    int k = threadIdx.x;
+    for (int i = k + (blockIdx.x * ROWS_PER_CTA); (k < ROWS_PER_CTA) && (i < N_ROWS);
+         k += ROWS_PER_CTA, i += ROWS_PER_CTA)
+    {
+      b_shared[i % (ROWS_PER_CTA + 1)] = b[i];
+    }
+  }
+
+  cg::sync(cta);
+
+  cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);
+
+  for (int k = 0, i = blockIdx.x * ROWS_PER_CTA; (k < ROWS_PER_CTA) && (i < N_ROWS); k++, i++)
+  {
+    double rowThreadSum = 0.0;
+    for (int j = threadIdx.x; j < N_ROWS; j += blockDim.x)
+    {
+      rowThreadSum += (A[i * N_ROWS + j] * x_shared[j]);
+    }
+
+    for (int offset = tile32.size() / 2; offset > 0; offset /= 2)
+    {
+      rowThreadSum += tile32.shfl_down(rowThreadSum, offset);
+    }
+
+    if (tile32.thread_rank() == 0)
+    {
+      atomicAdd(&b_shared[i % (ROWS_PER_CTA + 1)], -rowThreadSum);
+    }
+  }
+
+  cg::sync(cta);
+
+  if (threadIdx.x < ROWS_PER_CTA)
+  {
+    cg::thread_block_tile<ROWS_PER_CTA> tile8 = cg::tiled_partition<ROWS_PER_CTA>(cta);
+    double temp_sum                           = 0.0;
+
+    int k = threadIdx.x;
+
+    for (int i = k + (blockIdx.x * ROWS_PER_CTA); (k < ROWS_PER_CTA) && (i < N_ROWS);
+         k += ROWS_PER_CTA, i += ROWS_PER_CTA)
+    {
+      double dx = b_shared[i % (ROWS_PER_CTA + 1)];
+      dx /= A[i * N_ROWS + i];
+
+      x_new[i] = (x_shared[i] + dx);
+      temp_sum += fabs(dx);
+    }
+
+    for (int offset = tile8.size() / 2; offset > 0; offset /= 2)
+    {
+      temp_sum += tile8.shfl_down(temp_sum, offset);
+    }
+
+    if (tile8.thread_rank() == 0)
+    {
+      atomicAdd(sum, temp_sum);
+    }
+  }
+}
+
+// Thread block size for finalError kernel should be multiple of 32
+static __global__ void finalError(const double* x, double* g_sum)
+{
+  // Handle to thread block group
+  cg::thread_block cta = cg::this_thread_block();
+  extern __shared__ double warpSum[];
+  double sum = 0.0;
+
+  int globalThreadId = blockIdx.x * blockDim.x + threadIdx.x;
+
+  for (int i = globalThreadId; i < N_ROWS; i += blockDim.x * gridDim.x)
+  {
+    double d = x[i] - 1.0;
+    sum += fabs(d);
+  }
+
+  cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);
+
+  for (int offset = tile32.size() / 2; offset > 0; offset /= 2)
+  {
+    sum += tile32.shfl_down(sum, offset);
+  }
+
+  if (tile32.thread_rank() == 0)
+  {
+    warpSum[threadIdx.x / warpSize] = sum;
+  }
+
+  cg::sync(cta);
+
+  double blockSum = 0.0;
+  if (threadIdx.x < (blockDim.x / warpSize))
+  {
+    blockSum = warpSum[threadIdx.x];
+  }
+
+  if (threadIdx.x < 32)
+  {
+    for (int offset = tile32.size() / 2; offset > 0; offset /= 2)
+    {
+      blockSum += tile32.shfl_down(blockSum, offset);
+    }
+    if (tile32.thread_rank() == 0)
+    {
+      atomicAdd(g_sum, blockSum);
+    }
+  }
+}
+
+// Run the Jacobi method for A*x = b on GPU without CUDA Graph.
+// double JacobiMethodGpu(const float *A, const double *b,
+//                       const float conv_threshold, const int max_iter,
+//                       double *x, double *x_new) {
+double JacobiMethodGpu(
+  stream_ctx& ctx,
+  logical_data<slice<float>>& A_handle,
+  logical_data<slice<double>>& b_handle,
+  const float conv_threshold,
+  const int max_iter,
+  logical_data<slice<double>>& x_handle,
+  logical_data<slice<double>>& x_new_handle)
+{
+  // CTA size
+  dim3 nthreads(256, 1, 1);
+  // grid size
+  dim3 nblocks((N_ROWS / ROWS_PER_CTA) + 2, 1, 1);
+
+  auto sum_handle = ctx.logical_data(shape_of<scalar_view<double>>()).set_symbol("sum");
+
+  int k;
+  for (k = 0; k < max_iter; k++)
+  {
+    auto x_mode     = (k & 1) == 0 ? access_mode::read : access_mode::rw;
+    auto x_new_mode = (k & 1) == 0 ? access_mode::rw : access_mode::read;
+
+    ctx.task(A_handle.read(),
+             b_handle.read(),
+             task_dep<slice<double>>(x_handle, x_mode),
+             task_dep<slice<double>>(x_new_handle, x_new_mode),
+             sum_handle.write())
+        .set_symbol("JacobiMethod")
+        ->*[&](cudaStream_t stream, auto A, auto b, auto x, auto x_new, auto d_sum) {
+              cuda_try(cudaMemsetAsync(d_sum.addr, 0, sizeof(double), stream));
+
+              if ((k & 1) == 0)
+              {
+                JacobiMethod<<<nblocks, nthreads, 0, stream>>>(
+                  A.data_handle(), b.data_handle(), conv_threshold, x.data_handle(), x_new.data_handle(), d_sum.addr);
+              }
+              else
+              {
+                JacobiMethod<<<nblocks, nthreads, 0, stream>>>(
+                  A.data_handle(), b.data_handle(), conv_threshold, x_new.data_handle(), x.data_handle(), d_sum.addr);
+              }
+            };
+
+    if (ctx.wait(sum_handle) <= conv_threshold)
+    {
+      break;
+    }
+  }
+
+  auto final_x_handle = ((k & 1) == 0) ? &x_new_handle : &x_handle;
+  ctx.task(sum_handle.write(), final_x_handle->read()).set_symbol("finalError")
+      ->*[&](cudaStream_t stream, auto d_sum, auto final_x) {
+            cuda_try(cudaMemsetAsync(d_sum.addr, 0, sizeof(double), stream));
+
+            nblocks.x            = (N_ROWS / nthreads.x) + 1;
+            size_t sharedMemSize = ((nthreads.x / 32) + 1) * sizeof(double);
+            finalError<<<nblocks, nthreads, sharedMemSize, stream>>>(final_x.data_handle(), d_sum.addr);
+          };
+
+  return ctx.wait(sum_handle);
+}
+
+// Run the Jacobi method for A*x = b on CPU.
+void JacobiMethodCPU(float* A, double* b, float conv_threshold, int max_iter, int* num_iter, double* x)
+{
+  double* x_new = (double*) calloc(N_ROWS, sizeof(double));
+  SCOPE(exit)
+  {
+    free(x_new);
+  };
+
+  int k = 0;
+
+  for (; k < max_iter; k++)
+  {
+    double sum = 0.0;
+    for (int i = 0; i < N_ROWS; i++)
+    {
+      double temp_dx = b[i];
+      for (int j = 0; j < N_ROWS; j++)
+      {
+        temp_dx -= A[i * N_ROWS + j] * x[j];
+      }
+      temp_dx /= A[i * N_ROWS + i];
+      x_new[i] += temp_dx;
+      sum += fabs(temp_dx);
+    }
+
+    for (int i = 0; i < N_ROWS; i++)
+    {
+      x[i] = x_new[i];
+    }
+
+    if (sum <= conv_threshold)
+    {
+      break;
+    }
+  }
+
+  *num_iter = k + 1;
+}
+
+template <typename Ctx>
+int run()
+{
+  Ctx ctx;
+
+  double* b = cuda_try<cudaMallocHost<double>>(N_ROWS * sizeof(double), 0);
+  SCOPE(exit)
+  {
+    cuda_try(cudaFreeHost(b));
+  };
+
+  float* A = cuda_try<cudaMallocHost<float>>(N_ROWS * N_ROWS * sizeof(float), 0);
+  SCOPE(exit)
+  {
+    cuda_try(cudaFreeHost(A));
+  };
+
+  memset(b, 0, N_ROWS * sizeof(double));
+
+  memset(A, 0, N_ROWS * N_ROWS * sizeof(float));
+
+  createLinearSystem(A, b);
+  // start with array of all zeroes
+  double* x = (double*) calloc(N_ROWS, sizeof(double));
+  SCOPE(exit)
+  {
+    free(x);
+  };
+
+  auto A_handle     = ctx.logical_data(A, N_ROWS * N_ROWS).set_symbol("A");
+  auto b_handle     = ctx.logical_data(b, N_ROWS).set_symbol("b");
+  auto x_handle     = ctx.logical_data(x, N_ROWS).set_symbol("x");
+  auto x_new_handle = ctx.logical_data(shape_of<slice<double>>(N_ROWS)).set_symbol("x_new");
+
+  float conv_threshold = 1.0e-2;
+  int max_iter         = 4 * N_ROWS * N_ROWS;
+  int cnt              = 0;
+
+  JacobiMethodCPU(A, b, conv_threshold, max_iter, &cnt, x);
+
+  double sum = 0.0;
+  // Compute error
+  for (int i = 0; i < N_ROWS; i++)
+  {
+    double d = x[i] - 1.0;
+    sum += fabs(d);
+  }
+
+  ctx.task(x_handle.write()).set_symbol("memset x")->*[&](cudaStream_t stream, auto d_x) {
+    cuda_try(cudaMemsetAsync(d_x.data_handle(), 0, sizeof(double) * N_ROWS, stream));
+  };
+
+  ctx.task(x_new_handle.write()).set_symbol("memset x_new")->*[](cudaStream_t stream, auto d_x_new) {
+    cuda_try(cudaMemsetAsync(d_x_new.data_handle(), 0, sizeof(double) * N_ROWS, stream));
+  };
+
+  double sumGPU = JacobiMethodGpu(ctx, A_handle, b_handle, conv_threshold, max_iter, x_handle, x_new_handle);
+
+  ctx.finalize();
+
+  if (fabs(sum - sumGPU) > conv_threshold)
+  {
+    printf("&&&& jacobiCudaGraphs FAILED\n");
+    return EXIT_FAILURE;
+  }
+  return 0;
+}
+
+int main()
+{
+  return run<stream_ctx>();
+  // run<graph_ctx>();
+}

cccl_upstream/cudax/test/stf/cuda-samples/4_CUDA_Libraries/conjugateGradientMultiDeviceCG/conjugateGradientMultiDeviceCG_custf.cu

@@ -0,0 +1,521 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ *  * Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ *  * Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in the
+ *    documentation and/or other materials provided with the distribution.
+ *  * Neither the name of NVIDIA CORPORATION nor the names of its
+ *    contributors may be used to endorse or promote products derived
+ *    from this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+ * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+ * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+ * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+ * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+ * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+ * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+ * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+/*
+ * This sample implements a conjugate gradient solver on multiple GPU using
+ * Unified Memory optimized prefetching and usage hints.
+ *
+ */
+
+// includes, system
+#include <iostream>
+#include <map>
+#include <set>
+#include <utility>
+
+#include <cuda_runtime.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+// Utilities and system includes
+#include <cuda/experimental/__places/partitions/blocked_partition.cuh>
+#include <cuda/experimental/stf.cuh>
+
+#include <cooperative_groups.h>
+
+#include <cooperative_groups/reduce.h>
+
+using namespace cuda::experimental::stf;
+
+namespace cg = cooperative_groups;
+
+const char* sSDKname = "conjugateGradientMultiDeviceCG";
+
+#define ENABLE_CPU_DEBUG_CODE 0
+#define THREADS_PER_BLOCK     64
+
+__device__ double grid_dot_result = 0.0;
+
+/* genTridiag: generate a random tridiagonal symmetric matrix */
+void genTridiag(slice<int> I, slice<int> J, slice<float> val, int N, int nz)
+{
+  I(0) = 0, J(0) = 0, J(1) = 1;
+  val(0) = (float) rand() / RAND_MAX + 10.0f;
+  val(1) = (float) rand() / RAND_MAX;
+  int start;
+
+  for (int i = 1; i < N; i++)
+  {
+    if (i > 1)
+    {
+      I(i) = I(i - 1) + 3;
+    }
+    else
+    {
+      I(1) = 2;
+    }
+
+    start        = (i - 1) * 3 + 2;
+    J(start)     = i - 1;
+    J(start + 1) = i;
+
+    if (i < N - 1)
+    {
+      J(start + 2) = i + 1;
+    }
+
+    val(start)     = val(start - 1);
+    val(start + 1) = (float) rand() / RAND_MAX + 10.0f;
+
+    if (i < N - 1)
+    {
+      val(start + 2) = (float) rand() / RAND_MAX;
+    }
+  }
+
+  I(N) = nz;
+}
+
+// I - contains location of the given non-zero element in the row of the matrix
+// J - contains location of the given non-zero element in the column of the
+// matrix val - contains values of the given non-zero elements of the matrix
+// inputVecX - input vector to be multiplied
+// outputVecY - resultant vector
+void cpuSpMV(int* I, int* J, float* val, int /*unused*/, int num_rows, float alpha, float* inputVecX, float* outputVecY)
+{
+  for (int i = 0; i < num_rows; i++)
+  {
+    int num_elems_this_row = I[i + 1] - I[i];
+
+    float output = 0.0;
+    for (int j = 0; j < num_elems_this_row; j++)
+    {
+      output += alpha * val[I[i] + j] * inputVecX[J[I[i] + j]];
+    }
+    outputVecY[i] = output;
+  }
+
+  return;
+}
+
+float dotProduct(float* vecA, float* vecB, int size)
+{
+  float result = 0.0;
+
+  for (int i = 0; i < size; i++)
+  {
+    result = result + (vecA[i] * vecB[i]);
+  }
+
+  return result;
+}
+
+void scaleVector(float* vec, float alpha, int size)
+{
+  for (int i = 0; i < size; i++)
+  {
+    vec[i] = alpha * vec[i];
+  }
+}
+
+void saxpy(float* x, float* y, float a, int size)
+{
+  for (int i = 0; i < size; i++)
+  {
+    y[i] = a * x[i] + y[i];
+  }
+}
+
+void cpuConjugateGrad(int* I, int* J, float* val, float* x, float* Ax, float* p, float* r, int nnz, int N, float tol)
+{
+  int max_iter = 10000;
+
+  float alpha   = 1.0;
+  float alpham1 = -1.0;
+  float r0      = 0.0, b, a, na;
+
+  cpuSpMV(I, J, val, nnz, N, alpha, x, Ax);
+  saxpy(Ax, r, alpham1, N);
+
+  float r1 = dotProduct(r, r, N);
+
+  int k = 1;
+
+  while (r1 > tol * tol && k <= max_iter)
+  {
+    if (k > 1)
+    {
+      b = r1 / r0;
+      scaleVector(p, b, N);
+
+      saxpy(r, p, alpha, N);
+    }
+    else
+    {
+      for (int i = 0; i < N; i++)
+      {
+        p[i] = r[i];
+      }
+    }
+
+    cpuSpMV(I, J, val, nnz, N, alpha, p, Ax);
+
+    float dot = dotProduct(p, Ax, N);
+    a         = r1 / dot;
+
+    saxpy(p, x, a, N);
+    na = -a;
+    saxpy(Ax, r, na, N);
+
+    r0 = r1;
+    r1 = dotProduct(r, r, N);
+
+    printf("\nCPU code iteration = %3d, residual = %e\n", k, sqrt(r1));
+    k++;
+  }
+}
+
+template <typename thread_hierarchy_t>
+__device__ void gpuSpMV(
+  slice<const int> I,
+  slice<const int> J,
+  slice<const float> val,
+  int nnz,
+  int num_rows,
+  float alpha,
+  slice<float> inputVecX,
+  slice<float> outputVecY,
+  const thread_hierarchy_t& t)
+{
+  for (int i = t.rank(); i < num_rows; i += t.size())
+  {
+    int row_elem           = I(i);
+    int next_row_elem      = I(i + 1);
+    int num_elems_this_row = next_row_elem - row_elem;
+
+    float output = 0.0;
+    for (int j = 0; j < num_elems_this_row; j++)
+    {
+      output += alpha * val(row_elem + j) * inputVecX(J(row_elem + j));
+    }
+
+    outputVecY(i) = output;
+  }
+}
+
+template <typename thread_hierarchy_t>
+__device__ void gpuSaxpy(slice<float> x, slice<float> y, float a, int size, const thread_hierarchy_t& t)
+{
+  for (int i = t.rank(); i < size; i += t.size())
+  {
+    y(i) = a * x(i) + y(i);
+  }
+}
+
+template <typename thread_hierarchy_t>
+__device__ double
+gpuDotProduct(slice<float> vecA, slice<float> vecB, int size, double* dot_result, thread_hierarchy_t& t)
+{
+  slice<double> tmp = t.template storage<double>(1);
+
+  cg::thread_block cta = cooperative_groups::this_thread_block();
+
+  double temp_sum = 0.0;
+
+  for (int i = t.rank(); i < size; i += t.size())
+  {
+    temp_sum += (double) (vecA(i) * vecB(i));
+  }
+  cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);
+  temp_sum                         = cg::reduce(tile32, temp_sum, cg::plus<double>());
+
+  if (tile32.thread_rank() == 0)
+  {
+    tmp[tile32.meta_group_rank()] = temp_sum;
+  }
+
+  cta.sync();
+
+  if (tile32.meta_group_rank() == 0)
+  {
+    temp_sum = tile32.thread_rank() < tile32.meta_group_size() ? tmp[tile32.thread_rank()] : 0.0;
+    temp_sum = cg::reduce(tile32, temp_sum, cg::plus<double>());
+
+    if (tile32.thread_rank() == 0)
+    {
+      atomicAdd(&grid_dot_result, temp_sum);
+    }
+  }
+
+  t.sync();
+
+  if (t.rank(0, -1) == 0)
+  {
+    atomicAdd_system(dot_result, grid_dot_result);
+    grid_dot_result = 0.0;
+  }
+
+  t.sync();
+  return *dot_result;
+}
+
+template <typename thread_hierarchy_t>
+__device__ void gpuCopyVector(slice<float> srcA, slice<float> destB, int size, const thread_hierarchy_t& t)
+{
+  for (int i = t.rank(); i < size; i += t.size())
+  {
+    destB(i) = srcA(i);
+  }
+}
+
+template <typename thread_hierarchy_t>
+__device__ void
+gpuScaleVectorAndSaxpy(slice<float> x, slice<float> y, float a, float scale, int size, const thread_hierarchy_t& t)
+{
+  for (int i = t.rank(); i < size; i += t.size())
+  {
+    y(i) = a * x(i) + scale * y(i);
+  }
+}
+
+template <typename thread_hierarchy_t>
+__device__ void multiGpuConjugateGradient(
+  thread_hierarchy_t t,
+  slice<const int> I,
+  slice<const int> J,
+  slice<const float> val,
+  slice<float> x,
+  slice<float> Ax,
+  slice<float> p,
+  slice<float> r,
+  double* dot_result,
+  int nnz,
+  int N,
+  float tol)
+{
+  const int max_iter = 10000;
+
+  float alpha   = 1.0;
+  float alpham1 = -1.0;
+  float r0      = 0.0, r1, b, a, na;
+
+  for (int i = t.rank(); i < N; i += t.size())
+  {
+    r[i] = 1.0;
+    x[i] = 0.0;
+  }
+
+  gpuSpMV(I, J, val, nnz, N, alpha, x, Ax, t);
+
+  gpuSaxpy(Ax, r, alpham1, N, t);
+
+  r1 = gpuDotProduct(r, r, N, dot_result, t);
+
+  int k = 1;
+  while (r1 > tol * tol && k <= max_iter)
+  {
+    if (k > 1)
+    {
+      b = r1 / r0;
+      gpuScaleVectorAndSaxpy(r, p, alpha, b, N, t);
+    }
+    else
+    {
+      gpuCopyVector(r, p, N, t);
+    }
+
+    gpuSpMV(I, J, val, nnz, N, alpha, p, Ax, t);
+
+    if (t.rank() == 0)
+    {
+      *dot_result = 0.0;
+    }
+
+    a = r1 / gpuDotProduct(p, Ax, N, dot_result, t);
+
+    gpuSaxpy(p, x, a, N, t);
+
+    na = -a;
+
+    gpuSaxpy(Ax, r, na, N, t);
+
+    r0 = r1;
+
+    if (t.rank() == 0)
+    {
+      *dot_result = 0.0;
+    }
+
+    r1 = gpuDotProduct(r, r, N, dot_result, t);
+
+    k++;
+  }
+}
+
+int main()
+{
+  stream_ctx ctx;
+#if 0
+    constexpr size_t kNumGpusRequired = 8;
+#else
+  constexpr size_t kNumGpusRequired = 1;
+#endif
+  int N = 0, nz = 0, *I = NULL, *J = NULL;
+  float* val      = NULL;
+  const float tol = 1e-5f;
+  float* x;
+  float rhs = 1.0;
+  float r1;
+  float *r, *p, *Ax;
+
+  // printf("Starting [%s]...\n", sSDKname);
+
+  /* Generate a random tridiagonal symmetric matrix in CSR format */
+  N  = 10485760 * 2;
+  nz = (N - 2) * 3 + 4;
+
+  I              = (int*) malloc(sizeof(int) * (N + 1));
+  J              = (int*) malloc(sizeof(int) * nz);
+  val            = (float*) malloc(sizeof(float) * nz);
+  float* val_cpu = (float*) malloc(sizeof(float) * nz);
+
+  auto handle_I   = ctx.logical_data(I, {(unsigned) (N + 1)});
+  auto handle_J   = ctx.logical_data(J, {(unsigned) nz});
+  auto handle_val = ctx.logical_data(val, {(unsigned) nz});
+
+  ctx.host_launch(handle_I.write(), handle_J.write(), handle_val.write())->*[=](auto I, auto J, auto val) {
+    genTridiag(I, J, val, N, nz);
+    memcpy(val_cpu, val.data_handle(), sizeof(float) * nz);
+  };
+
+  double* dot_result = (double*) malloc(sizeof(double));
+  dot_result[0]      = 0.0;
+
+  x  = (float*) malloc(sizeof(float) * N);
+  r  = (float*) malloc(sizeof(float) * N);
+  p  = (float*) malloc(sizeof(float) * N);
+  Ax = (float*) malloc(sizeof(float) * N);
+
+  auto handle_r          = ctx.logical_data(r, {(unsigned) N});
+  auto handle_p          = ctx.logical_data(p, {(unsigned) N});
+  auto handle_Ax         = ctx.logical_data(Ax, {(unsigned) N});
+  auto handle_x          = ctx.logical_data(x, {(unsigned) N});
+  auto handle_dot_result = ctx.logical_data(dot_result, {(unsigned) 1});
+
+  // std::cout << "\nRunning on GPUs = " << kNumGpusRequired << '\n';
+  const int sMemSize = sizeof(double) * ((THREADS_PER_BLOCK / 32) + 1);
+
+  // auto all_devs = exec_place::repeat<blocked_partition>(exec_place::device(0), kNumGpusRequired);
+  auto all_devs = exec_place::n_devices(kNumGpusRequired);
+
+  /* The grid size is 0 and will be computed upon launch */
+  auto spec = con(con(THREADS_PER_BLOCK, mem(sMemSize)));
+  ctx.launch(
+    spec,
+    all_devs,
+    handle_I.read(),
+    handle_J.read(),
+    handle_val.read(),
+    handle_x.write(),
+    handle_Ax.write(),
+    handle_p.write(),
+    handle_r.write(),
+    handle_dot_result.write())
+      ->*[=]
+    _CCCL_DEVICE(auto t,
+                 slice<const int> I,
+                 slice<const int> J,
+                 slice<const float> val,
+                 slice<float> x,
+                 slice<float> Ax,
+                 slice<float> p,
+                 slice<float> r,
+                 slice<double> dot_result) {
+      multiGpuConjugateGradient(t, I, J, val, x, Ax, p, r, dot_result.data_handle(), nz, N, tol);
+    };
+
+  ctx.finalize();
+
+  r1 = dot_result[0];
+
+  printf("GPU Final, residual = %e \n  ", sqrt(r1));
+
+#if ENABLE_CPU_DEBUG_CODE
+  float* Ax_cpu = (float*) malloc(sizeof(float) * N);
+  float* r_cpu  = (float*) malloc(sizeof(float) * N);
+  float* p_cpu  = (float*) malloc(sizeof(float) * N);
+  float* x_cpu  = (float*) malloc(sizeof(float) * N);
+
+  for (int i = 0; i < N; i++)
+  {
+    r_cpu[i]  = 1.0;
+    Ax_cpu[i] = x_cpu[i] = 0.0;
+  }
+  cpuConjugateGrad(I, J, val, x_cpu, Ax_cpu, p_cpu, r_cpu, nz, N, tol);
+#endif
+
+  float rsum, diff, err = 0.0;
+
+  for (int i = 0; i < N; i++)
+  {
+    rsum = 0.0;
+
+    for (int j = I[i]; j < I[i + 1]; j++)
+    {
+      rsum += val_cpu[j] * x[J[j]];
+    }
+
+    diff = fabs(rsum - rhs);
+
+    if (diff > err)
+    {
+      err = diff;
+    }
+  }
+
+#if ENABLE_CPU_DEBUG_CODE
+  free(Ax_cpu);
+  free(r_cpu);
+  free(p_cpu);
+  free(x_cpu);
+#endif
+
+  printf("Test Summary:  Error amount = %f \n", err);
+  fprintf(stdout, "&&&& conjugateGradientMultiDeviceCG %s\n", (sqrt(r1) < tol) ? "PASSED" : "FAILED");
+  exit((sqrt(r1) < tol) ? EXIT_SUCCESS : EXIT_FAILURE);
+}

cccl_upstream/cudax/test/stf/cuda-samples/5_Domain_Specific/MonteCarloMultiGPU_cudastf/MonteCarloMultiGPU.cu

@@ -0,0 +1,210 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ *  * Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ *  * Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in the
+ *    documentation and/or other materials provided with the distribution.
+ *  * Neither the name of NVIDIA CORPORATION nor the names of its
+ *    contributors may be used to endorse or promote products derived
+ *    from this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+ * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+ * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+ * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+ * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+ * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+ * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+ * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+/*
+ * This sample evaluates fair call price for a
+ * given set of European options using Monte Carlo approach.
+ * See supplied whitepaper for more explanations.
+ */
+
+#include "MonteCarlo_gold.cu"
+#include "MonteCarlo_kernel.cu"
+
+////////////////////////////////////////////////////////////////////////////////
+// Common functions
+////////////////////////////////////////////////////////////////////////////////
+float randFloat(float low, float high)
+{
+  float t = (float) rand() / (float) RAND_MAX;
+  return (1.0f - t) * low + t * high;
+}
+
+/// Utility function to tweak problem size for small GPUs
+int adjustProblemSize(int GPU_N, int default_nOptions)
+{
+  int nOptions = default_nOptions;
+
+  for (int i = 0; i < GPU_N; i++)
+  {
+    cudaDeviceProp deviceProp;
+    cuda_safe_call(cudaGetDeviceProperties(&deviceProp, i));
+    int cudaCores = 80;
+
+    if (cudaCores <= 32)
+    {
+      nOptions = (nOptions < cudaCores / 2 ? nOptions : cudaCores / 2);
+    }
+  }
+
+  return nOptions;
+}
+
+int adjustGridSize(int GPUIndex, int defaultGridSize)
+{
+  cudaDeviceProp deviceProp;
+  cuda_safe_call(cudaGetDeviceProperties(&deviceProp, GPUIndex));
+  int maxGridSize = deviceProp.multiProcessorCount * 40;
+  return ((defaultGridSize > maxGridSize) ? maxGridSize : defaultGridSize);
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// CPU reference functions
+///////////////////////////////////////////////////////////////////////////////
+extern "C" void MonteCarloCPU(TOptionValue& callValue, TOptionData optionData, float* h_Random, int pathN);
+
+// Black-Scholes formula for call options
+extern "C" void BlackScholesCall(float& CallResult, TOptionData optionData);
+
+////////////////////////////////////////////////////////////////////////////////
+// Single-threaded multi-GPU solver using STF
+////////////////////////////////////////////////////////////////////////////////
+static void multiSolver(TOptionPlan* plan, int nPlans)
+{
+  stream_ctx ctx;
+
+  for (int i = 0; i < nPlans; i++)
+  {
+    cuda_safe_call(cudaSetDevice(plan[i].device));
+
+    initMonteCarloGPU(ctx, &plan[i]);
+    MonteCarloGPU(ctx, &plan[i]);
+    closeMonteCarloGPU(ctx, &plan[i]);
+  }
+
+  ctx.finalize();
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// Main program
+///////////////////////////////////////////////////////////////////////////////
+int main(int argc, char** argv)
+{
+  int GPU_N;
+  cuda_safe_call(cudaGetDeviceCount(&GPU_N));
+  int nOptions = 8 * 1024;
+
+  nOptions = adjustProblemSize(GPU_N, nOptions);
+
+  int OPT_N  = nOptions * GPU_N;
+  int PATH_N = 262144;
+
+  // Input data array
+  TOptionData* optionData = new TOptionData[OPT_N];
+  // Final GPU MC results
+  TOptionValue* callValueGPU = new TOptionValue[OPT_N];
+  //"Theoretical" call values by Black-Scholes formula
+  float* callValueBS = new float[OPT_N];
+  // Solver config
+  TOptionPlan* optionSolver = new TOptionPlan[GPU_N];
+
+  int i;
+  double delta, ref, sumDelta, sumRef, sumReserve;
+
+  srand(123);
+
+  for (i = 0; i < OPT_N; i++)
+  {
+    optionData[i].S            = randFloat(5.0f, 50.0f);
+    optionData[i].X            = randFloat(10.0f, 25.0f);
+    optionData[i].T            = randFloat(1.0f, 5.0f);
+    optionData[i].R            = 0.06f;
+    optionData[i].V            = 0.10f;
+    callValueGPU[i].Expected   = -1.0f;
+    callValueGPU[i].Confidence = -1.0f;
+  }
+
+  // Get option count for each GPU
+  for (i = 0; i < GPU_N; i++)
+  {
+    optionSolver[i].optionCount = OPT_N / GPU_N;
+  }
+
+  // Take into account cases with "odd" option counts
+  for (i = 0; i < (OPT_N % GPU_N); i++)
+  {
+    optionSolver[i].optionCount++;
+  }
+
+  // Assign GPU option ranges
+  int gpuBase = 0;
+
+  for (i = 0; i < GPU_N; i++)
+  {
+    optionSolver[i].device     = i;
+    optionSolver[i].optionData = optionData + gpuBase;
+    optionSolver[i].callValue  = callValueGPU + gpuBase;
+    optionSolver[i].pathN      = PATH_N;
+    optionSolver[i].gridSize   = adjustGridSize(optionSolver[i].device, optionSolver[i].optionCount);
+    gpuBase += optionSolver[i].optionCount;
+  }
+
+  multiSolver(optionSolver, GPU_N);
+
+  // Compare Monte Carlo and Black-Scholes results
+  sumDelta   = 0;
+  sumRef     = 0;
+  sumReserve = 0;
+
+  for (i = 0; i < OPT_N; i++)
+  {
+    BlackScholesCall(callValueBS[i], optionData[i]);
+    delta = fabs(callValueBS[i] - callValueGPU[i].Expected);
+    ref   = callValueBS[i];
+    sumDelta += delta;
+    sumRef += fabs(ref);
+
+    if (delta > 1e-6)
+    {
+      sumReserve += callValueGPU[i].Confidence / delta;
+    }
+  }
+
+  sumReserve /= OPT_N;
+
+  delete[] optionSolver;
+  delete[] callValueBS;
+  delete[] callValueGPU;
+  delete[] optionData;
+
+  if (sumReserve <= 1.0f)
+  {
+    printf("Test failed!\n");
+    return EXIT_FAILURE;
+  }
+
+  return EXIT_SUCCESS;
+}

cccl_upstream/cudax/test/stf/cuda-samples/5_Domain_Specific/MonteCarloMultiGPU_cudastf/MonteCarlo_common.cuh

@@ -0,0 +1,109 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ *  * Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ *  * Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in the
+ *    documentation and/or other materials provided with the distribution.
+ *  * Neither the name of NVIDIA CORPORATION nor the names of its
+ *    contributors may be used to endorse or promote products derived
+ *    from this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+ * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+ * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+ * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+ * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+ * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+ * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+ * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+#ifndef MONTECARLO_COMMON_H
+#define MONTECARLO_COMMON_H
+#include <cuda/experimental/stf.cuh>
+
+#include "curand_kernel.h"
+#include "realtype.cuh"
+
+using namespace cuda::experimental::stf;
+
+////////////////////////////////////////////////////////////////////////////////
+// Global types
+////////////////////////////////////////////////////////////////////////////////
+typedef struct
+{
+  float S;
+  float X;
+  float T;
+  float R;
+  float V;
+} TOptionData;
+
+typedef struct
+{
+  float Expected;
+  float Confidence;
+} TOptionValue;
+
+// Preprocessed input option data
+typedef struct
+{
+  real S;
+  real X;
+  real MuByT;
+  real VBySqrtT;
+} __TOptionData;
+
+// GPU outputs before CPU postprocessing
+typedef struct
+{
+  real Expected;
+  real Confidence;
+} __TOptionValue;
+
+typedef struct
+{
+  // Device ID for multi-GPU version
+  int device;
+  // Option count for this plan
+  int optionCount;
+
+  // Host-side data source and result destination
+  TOptionData* optionData;
+  TOptionValue* callValue;
+  logical_data<slice<__TOptionData>> preproc_optionData_handle;
+  logical_data<slice<__TOptionValue>> callValue_handle;
+
+  // Temporary Host-side pinned memory for async + faster data transfers
+  __TOptionValue* h_CallValue;
+
+  // Host-side option data
+  void* h_OptionData;
+
+  // Random number generator states
+  logical_data<slice<curandState>> rngStates_handle;
+
+  // Pseudorandom samples count
+  int pathN;
+
+  int gridSize;
+} TOptionPlan;
+
+#endif

cccl_upstream/cudax/test/stf/cuda-samples/5_Domain_Specific/MonteCarloMultiGPU_cudastf/MonteCarlo_gold.cu

@@ -0,0 +1,160 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ *  * Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ *  * Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in the
+ *    documentation and/or other materials provided with the distribution.
+ *  * Neither the name of NVIDIA CORPORATION nor the names of its
+ *    contributors may be used to endorse or promote products derived
+ *    from this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+ * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+ * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+ * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+ * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+ * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+ * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+ * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+#include "MonteCarlo_common.cuh"
+
+////////////////////////////////////////////////////////////////////////////////
+// Black-Scholes formula for Monte Carlo results validation
+////////////////////////////////////////////////////////////////////////////////
+#define A1       0.31938153
+#define A2       -0.356563782
+#define A3       1.781477937
+#define A4       -1.821255978
+#define A5       1.330274429
+#define RSQRT2PI 0.39894228040143267793994605993438
+
+// Polynomial approximation of
+// cumulative normal distribution function
+double CND(double d)
+{
+  double K = 1.0 / (1.0 + 0.2316419 * fabs(d));
+
+  double cnd = RSQRT2PI * exp(-0.5 * d * d) * (K * (A1 + K * (A2 + K * (A3 + K * (A4 + K * A5)))));
+
+  if (d > 0)
+  {
+    cnd = 1.0 - cnd;
+  }
+
+  return cnd;
+}
+
+// Black-Scholes formula for call value
+extern "C" void BlackScholesCall(float& callValue, TOptionData optionData)
+{
+  double S = optionData.S;
+  double X = optionData.X;
+  double T = optionData.T;
+  double R = optionData.R;
+  double V = optionData.V;
+
+  double sqrtT = sqrt(T);
+  double d1    = (log(S / X) + (R + 0.5 * V * V) * T) / (V * sqrtT);
+  double d2    = d1 - V * sqrtT;
+  double CNDD1 = CND(d1);
+  double CNDD2 = CND(d2);
+  double expRT = exp(-R * T);
+
+  callValue = (float) (S * CNDD1 - X * expRT * CNDD2);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+// CPU Monte Carlo
+////////////////////////////////////////////////////////////////////////////////
+static double CPU_endCallValue(double S, double X, double r, double MuByT, double VBySqrtT)
+{
+  double callValue = S * exp(MuByT + VBySqrtT * r) - X;
+  return (callValue > 0) ? callValue : 0;
+}
+
+#define CURAND_SAFE_CALL(call)                                                                         \
+  do                                                                                                   \
+  {                                                                                                    \
+    const curandStatus_t err = (call);                                                                 \
+    if (CURAND_STATUS_SUCCESS != err)                                                                  \
+    {                                                                                                  \
+      int dev = -1;                                                                                    \
+      cudaGetDevice(&dev);                                                                             \
+      fprintf(stderr, "%s:%u [device %d] CURAND error in call %s.\n", __FILE__, __LINE__, dev, #call); \
+      abort();                                                                                         \
+    }                                                                                                  \
+  } while (0)
+
+extern "C" void MonteCarloCPU(TOptionValue& callValue, TOptionData optionData, float* h_Samples, int pathN)
+{
+  const double S        = optionData.S;
+  const double X        = optionData.X;
+  const double T        = optionData.T;
+  const double R        = optionData.R;
+  const double V        = optionData.V;
+  const double MuByT    = (R - 0.5 * V * V) * T;
+  const double VBySqrtT = V * sqrt(T);
+
+  float* samples;
+  curandGenerator_t gen;
+
+  CURAND_SAFE_CALL(curandCreateGeneratorHost(&gen, CURAND_RNG_PSEUDO_DEFAULT));
+  unsigned long long seed = 1234ULL;
+  CURAND_SAFE_CALL(curandSetPseudoRandomGeneratorSeed(gen, seed));
+
+  if (h_Samples != NULL)
+  {
+    samples = h_Samples;
+  }
+  else
+  {
+    samples = (float*) malloc(pathN * sizeof(float));
+    CURAND_SAFE_CALL(curandGenerateNormal(gen, samples, pathN, 0.0, 1.0));
+  }
+
+  // for(int i=0; i<10; i++) printf("CPU sample = %f\n", samples[i]);
+
+  double sum = 0, sum2 = 0;
+
+  for (int pos = 0; pos < pathN; pos++)
+  {
+    double sample    = samples[pos];
+    double callValue = CPU_endCallValue(S, X, sample, MuByT, VBySqrtT);
+    sum += callValue;
+    sum2 += callValue * callValue;
+  }
+
+  if (h_Samples == NULL)
+  {
+    free(samples);
+  }
+
+  CURAND_SAFE_CALL(curandDestroyGenerator(gen));
+
+  // Derive average from the total sum and discount by riskfree rate
+  callValue.Expected = (float) (exp(-R * T) * sum / (double) pathN);
+  // Standard deviation
+  double stdDev = sqrt(((double) pathN * sum2 - sum * sum) / ((double) pathN * (double) (pathN - 1)));
+  // Confidence width; in 95% of all cases theoretical value lies within these
+  // borders
+  callValue.Confidence = (float) (exp(-R * T) * 1.96 * stdDev / sqrt((double) pathN));
+}

cccl_upstream/cudax/test/stf/cuda-samples/5_Domain_Specific/MonteCarloMultiGPU_cudastf/MonteCarlo_kernel.cu

@@ -0,0 +1,246 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ *  * Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ *  * Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in the
+ *    documentation and/or other materials provided with the distribution.
+ *  * Neither the name of NVIDIA CORPORATION nor the names of its
+ *    contributors may be used to endorse or promote products derived
+ *    from this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+ * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+ * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+ * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+ * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+ * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+ * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+ * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+////////////////////////////////////////////////////////////////////////////////
+// Global types
+////////////////////////////////////////////////////////////////////////////////
+#include "MonteCarlo_reduction.cuh"
+
+// This will output the proper error string when calling cudaGetLastError
+#define getLastCudaError(msg) __getLastCudaError((msg), __FILE__, __LINE__)
+
+namespace cg = cooperative_groups;
+
+inline void __getLastCudaError(const char* errorMessage, const char* file, const int line)
+{
+  cudaError_t err = cudaGetLastError();
+
+  if (cudaSuccess != err)
+  {
+    fprintf(stderr,
+            "%s(%i) : getLastCudaError() CUDA error:"
+            " %s : (%d) %s.\n",
+            file,
+            line,
+            errorMessage,
+            static_cast<int>(err),
+            cudaGetErrorString(err));
+    exit(EXIT_FAILURE);
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////
+// Internal GPU-side data structures
+////////////////////////////////////////////////////////////////////////////////
+#define MAX_OPTIONS (1024 * 1024)
+
+////////////////////////////////////////////////////////////////////////////////
+// Overloaded shortcut payoff functions for different precision modes
+////////////////////////////////////////////////////////////////////////////////
+__device__ inline float endCallValue(float S, float X, float r, float MuByT, float VBySqrtT)
+{
+  float callValue = S * __expf(MuByT + VBySqrtT * r) - X;
+  return (callValue > 0.0F) ? callValue : 0.0F;
+}
+
+__device__ inline double endCallValue(double S, double X, double r, double MuByT, double VBySqrtT)
+{
+  double callValue = S * exp(MuByT + VBySqrtT * r) - X;
+  return (callValue > 0.0) ? callValue : 0.0;
+}
+
+#define THREAD_N 256
+
+////////////////////////////////////////////////////////////////////////////////
+// This kernel computes the integral over all paths using a single thread block
+// per option. It is fastest when the number of thread blocks times the work per
+// block is high enough to keep the GPU busy.
+////////////////////////////////////////////////////////////////////////////////
+static __global__ void MonteCarloOneBlockPerOption(
+  curandState* __restrict rngStates,
+  const __TOptionData* __restrict d_OptionData,
+  __TOptionValue* __restrict d_CallValue,
+  int pathN,
+  int optionN)
+{
+  // Handle to thread block group
+  cg::thread_block cta             = cg::this_thread_block();
+  cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);
+
+  const int SUM_N = THREAD_N;
+  __shared__ real s_SumCall[SUM_N];
+  __shared__ real s_Sum2Call[SUM_N];
+
+  // determine global thread id
+  int tid = threadIdx.x + blockIdx.x * blockDim.x;
+
+  // Copy random number state to local memory for efficiency
+  curandState localState = rngStates[tid];
+  for (int optionIndex = blockIdx.x; optionIndex < optionN; optionIndex += gridDim.x)
+  {
+    const real S        = d_OptionData[optionIndex].S;
+    const real X        = d_OptionData[optionIndex].X;
+    const real MuByT    = d_OptionData[optionIndex].MuByT;
+    const real VBySqrtT = d_OptionData[optionIndex].VBySqrtT;
+
+    // Cycle through the entire samples array:
+    // derive end stock price for each path
+    // accumulate partial integrals into intermediate shared memory buffer
+    for (int iSum = threadIdx.x; iSum < SUM_N; iSum += blockDim.x)
+    {
+      __TOptionValue sumCall = {0, 0};
+
+      _CCCL_PRAGMA_UNROLL(8)
+      for (int i = iSum; i < pathN; i += SUM_N)
+      {
+        real r         = curand_normal(&localState);
+        real callValue = endCallValue(S, X, r, MuByT, VBySqrtT);
+        sumCall.Expected += callValue;
+        sumCall.Confidence += callValue * callValue;
+      }
+
+      s_SumCall[iSum]  = sumCall.Expected;
+      s_Sum2Call[iSum] = sumCall.Confidence;
+    }
+
+    // Reduce shared memory accumulators
+    // and write final result to global memory
+    cg::sync(cta);
+    sumReduce<real, SUM_N, THREAD_N>(s_SumCall, s_Sum2Call, cta, tile32, &d_CallValue[optionIndex]);
+  }
+}
+
+static __global__ void rngSetupStates(curandState* rngState, int device_id)
+{
+  // determine global thread id
+  int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  // Each threadblock gets different seed,
+  // Threads within a threadblock get different sequence numbers
+  curand_init(blockIdx.x + gridDim.x * device_id, threadIdx.x, 0, &rngState[tid]);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+// Host-side interface to GPU Monte Carlo
+////////////////////////////////////////////////////////////////////////////////
+
+template <typename Ctx>
+void initMonteCarloGPU(Ctx& ctx, TOptionPlan* plan)
+{
+  plan->h_OptionData = new __TOptionData[plan->optionCount];
+  plan->h_CallValue  = new __TOptionValue[plan->optionCount];
+  cuda_safe_call(
+    cudaHostRegister(plan->h_OptionData, plan->optionCount * sizeof(__TOptionData), cudaHostRegisterPortable));
+  cuda_safe_call(
+    cudaHostRegister(plan->h_CallValue, plan->optionCount * sizeof(__TOptionValue), cudaHostRegisterPortable));
+
+  // Register this vector
+  plan->preproc_optionData_handle =
+    ctx.logical_data((__TOptionData*) plan->h_OptionData, plan->optionCount).set_symbol("preproc_optionData");
+  plan->callValue_handle =
+    ctx.logical_data((__TOptionValue*) plan->h_CallValue, plan->optionCount).set_symbol("callValue");
+  plan->rngStates_handle =
+    ctx.logical_data(shape_of<slice<curandState>>(plan->gridSize * THREAD_N)).set_symbol("rngStates");
+
+  cuda_safe_call(cudaSetDevice(plan->device));
+
+  // Allocate states for pseudo random number generators
+  ctx.task(plan->rngStates_handle.write()).set_symbol("rngSetupStates")->*[&](cudaStream_t stream, auto rngStates) {
+    cuda_safe_call(
+      cudaMemsetAsync(rngStates.data_handle(), 0, plan->gridSize * THREAD_N * sizeof(curandState), stream));
+    getLastCudaError("cudaMemsetAsync failed.\n");
+
+    rngSetupStates<<<plan->gridSize, THREAD_N, 0, stream>>>(rngStates.data_handle(), plan->device);
+    getLastCudaError("rngSetupStates kernel failed.\n");
+  };
+}
+
+// Compute statistics and deallocate internal device memory
+template <typename Ctx>
+void closeMonteCarloGPU(Ctx& ctx, TOptionPlan* plan)
+{
+  ctx.host_launch(plan->callValue_handle.rw()).set_symbol("compute_stats")->*[&](auto h_CallValue) {
+    for (int i = 0; i < plan->optionCount; i++)
+    {
+      const double RT               = plan->optionData[i].R * plan->optionData[i].T;
+      const double sum              = h_CallValue.data_handle()[i].Expected;
+      const double sum2             = h_CallValue.data_handle()[i].Confidence;
+      const double pathN            = plan->pathN;
+      plan->callValue[i].Expected   = (float) (exp(-RT) * sum / pathN);
+      double stdDev                 = sqrt((pathN * sum2 - sum * sum) / (pathN * (pathN - 1)));
+      plan->callValue[i].Confidence = (float) (exp(-RT) * 1.96 * stdDev / sqrt(pathN));
+    }
+  };
+}
+
+// Main computations
+template <typename Ctx>
+void MonteCarloGPU(Ctx& ctx, TOptionPlan* plan)
+{
+  if (plan->optionCount <= 0 || plan->optionCount > MAX_OPTIONS)
+  {
+    printf("MonteCarloGPU(): bad option count.\n");
+    return;
+  }
+
+  // Preprocess computations on the host
+  ctx.host_launch(plan->preproc_optionData_handle.rw()).set_symbol("preprocess")->*[&](auto h_preproc_OptionData) {
+    for (int i = 0; i < plan->optionCount; i++)
+    {
+      const double T                                 = plan->optionData[i].T;
+      const double R                                 = plan->optionData[i].R;
+      const double V                                 = plan->optionData[i].V;
+      const double MuByT                             = (R - 0.5 * V * V) * T;
+      const double VBySqrtT                          = V * sqrt(T);
+      h_preproc_OptionData.data_handle()[i].S        = (real) plan->optionData[i].S;
+      h_preproc_OptionData.data_handle()[i].X        = (real) plan->optionData[i].X;
+      h_preproc_OptionData.data_handle()[i].MuByT    = (real) MuByT;
+      h_preproc_OptionData.data_handle()[i].VBySqrtT = (real) VBySqrtT;
+    }
+  };
+
+  ctx.task(plan->preproc_optionData_handle.read(), plan->callValue_handle.write(), plan->rngStates_handle.rw())
+      .set_symbol("MonteCarloOneBlockPerOption")
+      ->*[&](cudaStream_t stream, auto preproc_optionData, auto callValue_handle, auto rngStates) {
+            MonteCarloOneBlockPerOption<<<plan->gridSize, THREAD_N, 0, stream>>>(
+              rngStates.data_handle(),
+              preproc_optionData.data_handle(),
+              callValue_handle.data_handle(),
+              plan->pathN,
+              plan->optionCount);
+            getLastCudaError("MonteCarloOneBlockPerOption() execution failed\n");
+          };
+}

cccl_upstream/cudax/test/stf/cuda-samples/5_Domain_Specific/MonteCarloMultiGPU_cudastf/MonteCarlo_reduction.cuh

@@ -0,0 +1,85 @@
+/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ *  * Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ *  * Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in the
+ *    documentation and/or other materials provided with the distribution.
+ *  * Neither the name of NVIDIA CORPORATION nor the names of its
+ *    contributors may be used to endorse or promote products derived
+ *    from this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+ * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+ * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+ * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+ * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+ * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+ * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+ * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+#ifndef MONTECARLO_REDUCTION_CUH
+#define MONTECARLO_REDUCTION_CUH
+
+#include <cooperative_groups.h>
+
+#include "MonteCarlo_common.cuh"
+
+namespace cg = cooperative_groups;
+
+////////////////////////////////////////////////////////////////////////////////
+// This function calculates total sum for each of the two input arrays.
+// SUM_N must be power of two
+// Unrolling provides a bit of a performance improvement for small
+// to medium path counts.
+////////////////////////////////////////////////////////////////////////////////
+
+template <class T, int SUM_N, int blockSize>
+__device__ void
+sumReduce(T* sum, T* sum2, cg::thread_block& cta, cg::thread_block_tile<32>& tile32, __TOptionValue* d_CallValue)
+{
+  const int VEC = 32;
+  const int tid = cta.thread_rank();
+
+  T beta  = sum[tid];
+  T beta2 = sum2[tid];
+  T temp, temp2;
+
+  for (int i = VEC / 2; i > 0; i >>= 1)
+  {
+    if (tile32.thread_rank() < i)
+    {
+      temp  = sum[tid + i];
+      temp2 = sum2[tid + i];
+      beta += temp;
+      beta2 += temp2;
+      sum[tid]  = beta;
+      sum2[tid] = beta2;
+    }
+    cg::sync(tile32);
+  }
+  cg::sync(cta);
+
+  if (tid == 0)
+  {
+    beta  = 0;
+    beta2 = 0;
+    for (int i = 0; i < blockDim.x; i += VEC)
+    {
+      beta += sum[i];
+      beta2 += sum2[i];
+    }
+    __TOptionValue t = {beta, beta2};
+    *d_CallValue     = t;
+  }
+  cg::sync(cta);
+}
+
+#endif

cccl_upstream/cudax/test/stf/cuda-samples/5_Domain_Specific/MonteCarloMultiGPU_cudastf/README.md

@@ -0,0 +1,75 @@
+# MonteCarloMultiGPU - Monte Carlo Option Pricing with Multi-GPU support
+
+## Description
+
+This sample evaluates fair call price for a given set of European options using the Monte Carlo approach, taking advantage of all CUDA-capable GPUs installed in the system. This sample use double precision hardware if a GTX 200 class GPU is present.  The sample also takes advantage of CUDA 4.0 capability to supporting using a single CPU thread to control multiple GPUs
+
+## Key Concepts
+
+Random Number Generator, Computational Finance, CURAND Library
+
+## Supported SM Architectures
+
+[SM 3.5 ](https://developer.nvidia.com/cuda-gpus)  [SM 3.7 ](https://developer.nvidia.com/cuda-gpus)  [SM 5.0 ](https://developer.nvidia.com/cuda-gpus)  [SM 5.2 ](https://developer.nvidia.com/cuda-gpus)  [SM 5.3 ](https://developer.nvidia.com/cuda-gpus)  [SM 6.0 ](https://developer.nvidia.com/cuda-gpus)  [SM 6.1 ](https://developer.nvidia.com/cuda-gpus)  [SM 7.0 ](https://developer.nvidia.com/cuda-gpus)  [SM 7.2 ](https://developer.nvidia.com/cuda-gpus)  [SM 7.5 ](https://developer.nvidia.com/cuda-gpus)  [SM 8.0 ](https://developer.nvidia.com/cuda-gpus)  [SM 8.6 ](https://developer.nvidia.com/cuda-gpus)  [SM 8.7 ](https://developer.nvidia.com/cuda-gpus)
+
+## Supported OSes
+
+Linux, Windows
+
+## Supported CPU Architecture
+
+x86_64, ppc64le, armv7l
+
+## CUDA APIs involved
+
+### [CUDA Runtime API](http://docs.nvidia.com/cuda/cuda-runtime-api/index.html)
+cudaMemset, cudaFree, cudaStreamDestroy, cudaEventRecord, cudaMallocHost, cudaStreamCreate, cudaEventCreate, cudaGetDeviceCount, cudaDeviceSynchronize, cudaEventSynchronize, cudaFreeHost, cudaMalloc, cudaEventDestroy, cudaSetDevice, cudaMemcpyAsync, cudaStreamSynchronize, cudaGetDeviceProperties
+
+## Dependencies needed to build/run
+[CURAND](../../../README.md#curand)
+
+## Prerequisites
+
+Download and install the [CUDA Toolkit 11.6](https://developer.nvidia.com/cuda-downloads) for your corresponding platform.
+Make sure the dependencies mentioned in [Dependencies]() section above are installed.
+
+## Build and Run
+
+### Windows
+The Windows samples are built using the Visual Studio IDE. Solution files (.sln) are provided for each supported version of Visual Studio, using the format:
+```
+*_vs<version>.sln - for Visual Studio <version>
+```
+Each individual sample has its own set of solution files in its directory:
+
+To build/examine all the samples at once, the complete solution files should be used. To build/examine a single sample, the individual sample solution files should be used.
+> **Note:** Some samples require that the Microsoft DirectX SDK (June 2010 or newer) be installed and that the VC++ directory paths are properly set up (**Tools > Options...**). Check DirectX Dependencies section for details."
+
+### Linux
+The Linux samples are built using makefiles. To use the makefiles, change the current directory to the sample directory you wish to build, and run make:
+```
+$ cd <sample_dir>
+$ make
+```
+The samples makefiles can take advantage of certain options:
+*  **TARGET_ARCH=<arch>** - cross-compile targeting a specific architecture. Allowed architectures are x86_64, ppc64le, armv7l.
+    By default, TARGET_ARCH is set to HOST_ARCH. On a x86_64 machine, not setting TARGET_ARCH is the equivalent of setting TARGET_ARCH=x86_64.<br/>
+`$ make TARGET_ARCH=x86_64` <br/> `$ make TARGET_ARCH=ppc64le` <br/> `$ make TARGET_ARCH=armv7l` <br/>
+    See [here](http://docs.nvidia.com/cuda/cuda-samples/index.html#cross-samples) for more details.
+*   **dbg=1** - build with debug symbols
+    ```
+    $ make dbg=1
+    ```
+*   **SMS="A B ..."** - override the SM architectures for which the sample will be built, where `"A B ..."` is a space-delimited list of SM architectures. For example, to generate SASS for SM 50 and SM 60, use `SMS="50 60"`.
+    ```
+    $ make SMS="50 60"
+    ```
+
+*  **HOST_COMPILER=<host_compiler>** - override the default g++ host compiler. See the [Linux Installation Guide](http://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html#system-requirements) for a list of supported host compilers.
+```
+    $ make HOST_COMPILER=g++
+```
+
+## References (for more details)
+
+[whitepaper](./doc/MonteCarlo.pdf)

cccl_upstream/cudax/test/stf/cuda-samples/5_Domain_Specific/MonteCarloMultiGPU_cudastf/realtype.cuh

@@ -0,0 +1,49 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ *  * Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ *  * Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in the
+ *    documentation and/or other materials provided with the distribution.
+ *  * Neither the name of NVIDIA CORPORATION nor the names of its
+ *    contributors may be used to endorse or promote products derived
+ *    from this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+ * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+ * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+ * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+ * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+ * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+ * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+ * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+#ifndef REALTYPE_H
+#define REALTYPE_H
+
+// #define DOUBLE_PRECISION
+
+#ifndef DOUBLE_PRECISION
+typedef float real;
+#else
+typedef double real;
+#endif
+
+#endif

cccl_upstream/cudax/test/stf/dot/basic.cu

@@ -0,0 +1,48 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief This test makes sure we can generate a dot file
+ */
+
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+int main()
+{
+// TODO (miscco): Make it work for windows
+#if !_CCCL_COMPILER(MSVC)
+  // Generate a random filename
+  int r = rand();
+
+  char filename[64];
+  snprintf(filename, 64, "output_%d.dot", r);
+  // fprintf(stderr, "filename %s\n", filename);
+  setenv("CUDASTF_DOT_FILE", filename, 1);
+
+  stream_ctx ctx;
+
+  auto lA = ctx.logical_data(shape_of<slice<char>>(64));
+  ctx.task(lA.write())->*[](cudaStream_t, auto) {};
+  ctx.task(lA.rw())->*[](cudaStream_t, auto) {};
+  ctx.finalize();
+
+  // Call this explicitly for the purpose of the test
+  reserved::dot::instance().finish();
+
+  // Make sure the file exists, and erase it
+  // fprintf(stderr, "ERASE. ...\n");
+  EXPECT(access(filename, F_OK) != -1);
+
+  EXPECT(unlink(filename) == 0);
+#endif // !_CCCL_COMPILER(MSVC)
+}

cccl_upstream/cudax/test/stf/dot/graph_print_to_dot.cu

@@ -0,0 +1,50 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief This test makes sure we can generate a dot file
+ */
+
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+__global__ void dummy() {}
+
+int main()
+{
+// TODO (miscco): Make it work for windows
+#if !_CCCL_COMPILER(MSVC)
+  // Generate a random filename
+  int r = rand();
+
+  char filename[64];
+  snprintf(filename, 64, "output_%d.dot", r);
+  // fprintf(stderr, "filename %s\n", filename);
+
+  graph_ctx ctx;
+
+  auto lA = ctx.logical_data(shape_of<slice<char>>(64));
+  ctx.task(lA.write())->*[](cudaStream_t s, auto) {
+    dummy<<<1, 1, 0, s>>>();
+  };
+  ctx.task(lA.rw())->*[](cudaStream_t s, auto) {
+    dummy<<<1, 1, 0, s>>>();
+  };
+  ctx.print_to_dot(filename, cudaGraphDebugDotFlagsVerbose);
+  ctx.finalize();
+
+  // Make sure the file exists, and erase it
+  EXPECT(access(filename, F_OK) != -1);
+
+  EXPECT(unlink(filename) == 0);
+#endif // !_CCCL_COMPILER(MSVC)
+}

cccl_upstream/cudax/test/stf/dot/section_movable.cu

@@ -0,0 +1,47 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2025 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief Test if dot section guards are movable and if we can them as an optional value
+ */
+
+#include <cuda/experimental/stf.cuh>
+
+#include <optional>
+#include <vector>
+
+using namespace cuda::experimental::stf;
+
+int main()
+{
+  context ctx;
+
+  // Ensure the guard is movable
+  {
+    auto g  = ctx.dot_section("foo");
+    auto g2 = mv(g);
+  }
+
+  // Ensure the guard can be stored as an optional
+  ::std::vector<::std::optional<reserved::dot_section::guard>> nested_sections;
+
+  for (size_t depth = 0; depth < 3; depth++)
+  {
+    nested_sections.emplace_back(ctx.dot_section("foo"));
+  }
+
+  for (size_t depth = 0; depth < 3; depth++)
+  {
+    nested_sections.pop_back();
+  }
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/dot/sections.cu

@@ -0,0 +1,63 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief This test makes sure we can generate a dot file with sections
+ */
+
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+int main()
+{
+// TODO (miscco): Make it work for windows
+#if !_CCCL_COMPILER(MSVC)
+  // Generate a random filename
+  int r = rand();
+
+  char filename[64];
+  snprintf(filename, 64, "output_%d.dot", r);
+  // fprintf(stderr, "filename %s\n", filename);
+  setenv("CUDASTF_DOT_FILE", filename, 1);
+
+  context ctx;
+
+  auto lA = ctx.logical_data(shape_of<slice<char>>(64));
+  auto lB = ctx.logical_data(shape_of<slice<char>>(64));
+  auto lC = ctx.logical_data(shape_of<slice<char>>(64));
+  ctx.task(lA.write()).set_symbol("initA")->*[](cudaStream_t, auto) {};
+  ctx.task(lB.write()).set_symbol("initB")->*[](cudaStream_t, auto) {};
+  ctx.task(lC.write()).set_symbol("initC")->*[](cudaStream_t, auto) {};
+  for (size_t j = 0; j < 3; j++)
+  {
+    ctx.task(lA.rw()).set_symbol("f1")->*[](cudaStream_t, auto) {};
+    auto guard = ctx.dot_section("sec_loop " + ::std::to_string(j));
+    for (size_t i = 0; i < 2; i++)
+    {
+      auto guard_inner = ctx.dot_section("sec_inner_loop " + ::std::to_string(i));
+      ctx.task(lA.read(), lB.rw()).set_symbol("f2")->*[](cudaStream_t, auto, auto) {};
+      ctx.task(lA.read(), lC.rw()).set_symbol("f2")->*[](cudaStream_t, auto, auto) {};
+      ctx.task(lB.read(), lC.read(), lA.rw()).set_symbol("f3")->*[](cudaStream_t, auto, auto, auto) {};
+    }
+  }
+  ctx.finalize();
+
+  // Call this explicitly for the purpose of the test
+  reserved::dot::instance().finish();
+
+  // Make sure the file exists, and erase it
+  // fprintf(stderr, "ERASE. ...\n");
+  EXPECT(access(filename, F_OK) != -1);
+
+  EXPECT(unlink(filename) == 0);
+#endif // !_CCCL_COMPILER(MSVC)
+}

cccl_upstream/cudax/test/stf/dot/sections_2.cu

@@ -0,0 +1,49 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief This test makes sure we can generate a dot file with sections
+ */
+
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+int main()
+{
+// TODO (miscco): Make it work for windows
+#if !_CCCL_COMPILER(MSVC)
+  context ctx;
+  auto lA = ctx.token().set_symbol("A");
+  auto lB = ctx.token().set_symbol("B");
+  auto lC = ctx.token().set_symbol("C");
+
+  // Begin a top-level section named "foo"
+  auto s_foo = ctx.dot_section("foo");
+  for (size_t i = 0; i < 2; i++)
+  {
+    // Section named "bar" using RAII
+    auto s_bar = ctx.dot_section("bar");
+    ctx.task(lA.read(), lB.rw()).set_symbol("t1")->*[](cudaStream_t) {};
+    for (size_t j = 0; j < 2; j++)
+    {
+      // Section named "baz" using RAII
+      auto s_bar = ctx.dot_section("baz");
+      ctx.task(lA.read(), lC.rw()).set_symbol("t2")->*[](cudaStream_t) {};
+      ctx.task(lB.read(), lC.read(), lA.rw()).set_symbol("t3")->*[](cudaStream_t) {};
+      // Implicit end of section "baz"
+    }
+    // Implicit end of section "bar"
+  }
+  s_foo.end(); // Explicit end of section "foo"
+  ctx.finalize();
+#endif // !_CCCL_COMPILER(MSVC)
+}

cccl_upstream/cudax/test/stf/dot/sections_stackable.cu

@@ -0,0 +1,61 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief This test makes sure we can generate a dot file with sections and stackable contexts
+ */
+
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+int main()
+{
+// TODO (miscco): Make it work for windows
+#if !_CCCL_COMPILER(MSVC)
+  // Test configuration constants
+  constexpr size_t buffer_size    = 64;
+  constexpr size_t num_iterations = 3;
+
+  stackable_ctx ctx;
+
+  // Create logical data for the computation pipeline
+  auto lA = ctx.logical_data(shape_of<slice<char>>(buffer_size));
+  auto lB = ctx.logical_data(shape_of<slice<char>>(buffer_size));
+  auto lC = ctx.logical_data(shape_of<slice<char>>(buffer_size));
+
+  // Initialize all logical data in a dedicated DOT section
+  auto r_init = ctx.dot_section("init");
+  ctx.task(lA.write()).set_symbol("initA")->*[](cudaStream_t, auto) {};
+  ctx.task(lB.write()).set_symbol("initB")->*[](cudaStream_t, auto) {};
+  ctx.task(lC.write()).set_symbol("initC")->*[](cudaStream_t, auto) {};
+  r_init.end();
+
+  // Test nested DOT sections with stackable contexts
+  // This creates a hierarchical structure to verify DOT graph generation
+  for (size_t j = 0; j < num_iterations; j++)
+  {
+    auto r0 = ctx.dot_section("lvl0");
+    ctx.task(lA.rw()).set_symbol("f1")->*[](cudaStream_t, auto) {};
+
+    ctx.push();
+    {
+      auto r1 = ctx.dot_section("lvl1");
+      ctx.task(lA.read(), lB.rw()).set_symbol("f2")->*[](cudaStream_t, auto, auto) {};
+      ctx.task(lA.read(), lC.rw()).set_symbol("f2")->*[](cudaStream_t, auto, auto) {};
+      ctx.task(lB.read(), lC.read(), lA.rw()).set_symbol("f3")->*[](cudaStream_t, auto, auto, auto) {};
+    }
+    ctx.pop();
+  }
+  ctx.finalize();
+
+#endif // !_CCCL_COMPILER(MSVC)
+}

cccl_upstream/cudax/test/stf/dot/with_events.cu

@@ -0,0 +1,49 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief This test makes sure we can generate a dot file with events
+ */
+
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+int main()
+{
+// TODO (miscco): Make it work for windows
+#if !_CCCL_COMPILER(MSVC)
+  // Generate a random filename
+  int r = rand();
+
+  char filename[64];
+  snprintf(filename, 64, "output_%d.dot", r);
+  // fprintf(stderr, "filename %s\n", filename);
+  setenv("CUDASTF_DOT_FILE", filename, 1);
+  setenv("CUDASTF_DOT_IGNORE_PREREQS", "0", 1);
+
+  stream_ctx ctx;
+
+  auto lA = ctx.logical_data(shape_of<slice<char>>(64));
+  ctx.task(lA.write())->*[](cudaStream_t, auto) {};
+  ctx.task(lA.rw())->*[](cudaStream_t, auto) {};
+  ctx.finalize();
+
+  // Call this explicitly for the purpose of the test
+  reserved::dot::instance().finish();
+
+  // Make sure the file exists, and erase it
+  // fprintf(stderr, "ERASE. ...\n");
+  EXPECT(access(filename, F_OK) != -1);
+
+  EXPECT(unlink(filename) == 0);
+#endif // !_CCCL_COMPILER(MSVC)
+}

cccl_upstream/cudax/test/stf/error_checks/ctx_mismatch.cu

@@ -0,0 +1,91 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief Ensure an error is detected when a task uses a logical data from a
+ *        different context
+ */
+
+#include <cuda/experimental/stf.cuh>
+
+#include <csignal>
+
+using namespace cuda::experimental::stf;
+
+bool should_abort = false;
+
+void cleanupRoutine(int /*unused*/)
+{
+  if (should_abort)
+  {
+    exit(EXIT_SUCCESS);
+  }
+  else
+  {
+    fprintf(stderr, "Unexpected SIGABRT !\n");
+    exit(EXIT_FAILURE);
+  }
+}
+
+template <typename Ctx, size_t n>
+void run(double (&X)[n])
+{
+  Ctx ctx1;
+  auto lX = ctx1.logical_data(X);
+
+  // We are now using lX in the wrong context
+  should_abort = true;
+
+  Ctx ctx2;
+  ctx2.task(lX.rw())->*[&](cudaStream_t /*unused*/, auto /*unused*/) {};
+
+  assert(0 && "This should not be reached");
+}
+
+int main()
+{
+  /* Setup an handler to catch the SIGABRT signal during the programming error */
+#if _CCCL_COMPILER(MSVC)
+  signal(SIGABRT, &cleanupRoutine);
+#else // ^^^ _CCCL_COMPILER(MSVC) ^^^ / vvv !_CCCL_COMPILER(MSVC)
+  struct sigaction sigabrt_action{};
+  memset(&sigabrt_action, 0, sizeof(sigabrt_action));
+  sigabrt_action.sa_handler = &cleanupRoutine;
+
+  if (sigaction(SIGABRT, &sigabrt_action, nullptr) != 0)
+  {
+    perror("sigaction SIGABRT");
+    exit(EXIT_FAILURE);
+  }
+#endif // !_CCCL_COMPILER(MSVC)
+
+  const int n = 12;
+  double X[n];
+
+  for (int ind = 0; ind < n; ind++)
+  {
+    X[ind] = 1.0 * ind;
+  }
+
+  // We can't run both stream and graph tests because either will abort the program. So choose one at random.
+  srand(static_cast<unsigned>(time(nullptr)));
+  if (rand() % 2 == 0)
+  {
+    run<stream_ctx>(X);
+  }
+  else
+  {
+    run<graph_ctx>(X);
+  }
+
+  assert(0 && "This should not be reached");
+  return EXIT_FAILURE;
+}

cccl_upstream/cudax/test/stf/error_checks/data_interface_mismatch.cu

@@ -0,0 +1,101 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief Ensure an error is detected dynamically if we access a data instance
+ *        with the wrong interface type
+ */
+
+#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+#include <csignal>
+
+using namespace cuda::experimental::stf;
+
+bool should_abort = false;
+
+void cleanupRoutine(int /*unused*/)
+{
+  if (should_abort)
+  {
+    exit(EXIT_SUCCESS);
+  }
+  else
+  {
+    fprintf(stderr, "Unexpected SIGABRT !\n");
+    exit(EXIT_FAILURE);
+  }
+}
+
+template <typename Ctx, size_t n>
+void run(double (&X)[n])
+{
+  Ctx ctx;
+  // This creates an untyped logical data that is implicitly a vector of size
+  // n. Had the code used `auto` instead of `logical_data_untyped`, errors
+  // would have been rejected statically. We want to disable static checking
+  // for the purposes of this test.
+  logical_data_untyped handle_X = ctx.logical_data(X);
+
+  // Here we create a dynamically-typed task, again to go around static typechecking.
+  auto t = ctx.task();
+  t.add_deps(handle_X.rw());
+
+  t->*[&](auto&) {
+    should_abort = true;
+    // We have a programming error here with a vector of `double` accessad as a vector of `float`.
+    handle_X.instance<slice<float>>(t);
+    should_abort = false;
+  };
+
+  assert(0 && "This should not be reached");
+}
+
+int main()
+{
+  /* Setup an handler to catch the SIGABRT signal during the programming error */
+#if _CCCL_COMPILER(MSVC)
+  signal(SIGABRT, &cleanupRoutine);
+#else // ^^^ _CCCL_COMPILER(MSVC) ^^^ / vvv !_CCCL_COMPILER(MSVC)
+  struct sigaction sigabrt_action{};
+  memset(&sigabrt_action, 0, sizeof(sigabrt_action));
+  sigabrt_action.sa_handler = &cleanupRoutine;
+
+  if (sigaction(SIGABRT, &sigabrt_action, nullptr) != 0)
+  {
+    perror("sigaction SIGABRT");
+    exit(EXIT_FAILURE);
+  }
+#endif // !_CCCL_COMPILER(MSVC)
+
+  const int n = 12;
+  double X[n];
+
+  for (int ind = 0; ind < n; ind++)
+  {
+    X[ind] = 1.0 * ind;
+  }
+
+  // We can't run both stream and graph tests because either will abort the program. So choose one at random.
+  srand(static_cast<unsigned>(time(nullptr)));
+  if (rand() % 2 == 0)
+  {
+    run<stream_ctx>(X);
+  }
+  else
+  {
+    run<graph_ctx>(X);
+  }
+
+  assert(0 && "This should not be reached");
+  return EXIT_FAILURE;
+}

cccl_upstream/cudax/test/stf/error_checks/double_finalize.cu

@@ -0,0 +1,80 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief Ensure an error is detected if we can finalize more than once
+ */
+
+#include <cuda/experimental/stf.cuh>
+
+#include <csignal>
+
+using namespace cuda::experimental::stf;
+
+bool should_abort = false;
+
+void cleanupRoutine(int /*unused*/)
+{
+  if (should_abort)
+  {
+    exit(EXIT_SUCCESS);
+  }
+  else
+  {
+    fprintf(stderr, "Unexpected SIGABRT !\n");
+    exit(EXIT_FAILURE);
+  }
+}
+
+int main()
+{
+  // This test only works when assert() is enabled in
+#ifndef NDEBUG
+  /* Setup an handler to catch the SIGABRT signal during the programming error */
+#  if _CCCL_COMPILER(MSVC)
+  signal(SIGABRT, &cleanupRoutine);
+#  else // ^^^ _CCCL_COMPILER(MSVC) ^^^ / vvv !_CCCL_COMPILER(MSVC)
+  struct sigaction sigabrt_action{};
+  memset(&sigabrt_action, 0, sizeof(sigabrt_action));
+  sigabrt_action.sa_handler = &cleanupRoutine;
+
+  if (sigaction(SIGABRT, &sigabrt_action, nullptr) != 0)
+  {
+    perror("sigaction SIGABRT");
+    exit(EXIT_FAILURE);
+  }
+#  endif // !_CCCL_COMPILER(MSVC)
+
+  context ctx;
+
+  const int n = 12;
+  double X[n];
+
+  for (int ind = 0; ind < n; ind++)
+  {
+    X[ind] = 1.0 * ind;
+  }
+
+  // This creates a handle that is implicitly a vector of size n
+  auto lX = ctx.logical_data(X);
+
+  ctx.task(lX.rw())->*[](cudaStream_t, auto) { /* no-op */ };
+
+  ctx.finalize();
+
+  should_abort = true;
+  // We cannot call sync twice
+  ctx.finalize();
+
+  assert(0 && "This should not be reached");
+  return EXIT_FAILURE;
+#endif
+}

cccl_upstream/cudax/test/stf/error_checks/erase_frozen.cu

@@ -0,0 +1,77 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ *
+ * @brief Ensure temporary data are destroyed
+ *
+ */
+
+#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+#include <csignal>
+
+using namespace cuda::experimental::stf;
+
+bool should_abort = false;
+
+void cleanupRoutine(int /*unused*/)
+{
+  if (should_abort)
+  {
+    exit(EXIT_SUCCESS);
+  }
+  else
+  {
+    fprintf(stderr, "Unexpected SIGABRT !\n");
+    exit(EXIT_FAILURE);
+  }
+}
+
+int main()
+{
+  /* Setup an handler to catch the SIGABRT signal during the programming error */
+#if _CCCL_COMPILER(MSVC)
+  signal(SIGABRT, &cleanupRoutine);
+#else // ^^^ _CCCL_COMPILER(MSVC) ^^^ / vvv !_CCCL_COMPILER(MSVC)
+  struct sigaction sigabrt_action{};
+  memset(&sigabrt_action, 0, sizeof(sigabrt_action));
+  sigabrt_action.sa_handler = &cleanupRoutine;
+
+  if (sigaction(SIGABRT, &sigabrt_action, nullptr) != 0)
+  {
+    perror("sigaction SIGABRT");
+    exit(EXIT_FAILURE);
+  }
+#endif // !_CCCL_COMPILER(MSVC)
+
+  stream_ctx ctx;
+  const int N = 16;
+  int X[N];
+
+  for (int i = 0; i < N; i++)
+  {
+    X[i] = i;
+  }
+
+  auto lX = ctx.logical_data(X);
+
+  lX.freeze(access_mode::rw, data_place::current_device());
+
+  // This should cause an error because lX is frozen while the context is finalized
+  should_abort = true;
+
+  ctx.finalize();
+
+  assert(0 && "This should not be reached");
+  return EXIT_FAILURE;
+}

cccl_upstream/cudax/test/stf/error_checks/misformed_tasks_dbl_end.cu

@@ -0,0 +1,78 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief Ensure that an error is detected if we end a task twice
+ */
+
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+#include <csignal>
+
+using namespace cuda::experimental::stf;
+
+bool should_abort = false;
+
+void cleanupRoutine(int /*unused*/)
+{
+  if (should_abort)
+  {
+    exit(EXIT_SUCCESS);
+  }
+  else
+  {
+    fprintf(stderr, "Unexpected SIGABRT !\n");
+    exit(EXIT_FAILURE);
+  }
+}
+
+int main()
+{
+  // This test only works when assert() is enabled in
+#ifndef NDEBUG
+  /* Setup an handler to catch the SIGABRT signal during the programming error */
+#  if _CCCL_COMPILER(MSVC)
+  signal(SIGABRT, &cleanupRoutine);
+#  else // ^^^ _CCCL_COMPILER(MSVC) ^^^ / vvv !_CCCL_COMPILER(MSVC)
+  struct sigaction sigabrt_action{};
+  memset(&sigabrt_action, 0, sizeof(sigabrt_action));
+  sigabrt_action.sa_handler = &cleanupRoutine;
+
+  if (sigaction(SIGABRT, &sigabrt_action, nullptr) != 0)
+  {
+    perror("sigaction SIGABRT");
+    exit(EXIT_FAILURE);
+  }
+#  endif // !_CCCL_COMPILER(MSVC)
+
+  stream_ctx ctx;
+
+  const int n = 12;
+  double X[n];
+
+  for (int ind = 0; ind < n; ind++)
+  {
+    X[ind] = 1.0 * ind;
+  }
+
+  // This creates a handle that is implicitly a vector of size n
+  auto lX = ctx.logical_data(X);
+
+  auto t = ctx.task(lX.rw());
+  t.start();
+  t.end();
+  should_abort = true;
+  t.end();
+
+  assert(0 && "This should not be reached");
+  return EXIT_FAILURE;
+#endif
+}

cccl_upstream/cudax/test/stf/error_checks/misformed_tasks_dbl_start.cu

@@ -0,0 +1,75 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief Ensure that an error is detected if we start a task twice
+ */
+
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+#include <csignal>
+
+using namespace cuda::experimental::stf;
+
+bool should_abort = false;
+
+void cleanupRoutine(int /*unused*/)
+{
+  if (should_abort)
+  {
+    exit(EXIT_SUCCESS);
+  }
+  else
+  {
+    fprintf(stderr, "Unexpected SIGABRT !\n");
+    exit(EXIT_FAILURE);
+  }
+}
+
+int main()
+{
+  /* Setup an handler to catch the SIGABRT signal during the programming error */
+#if _CCCL_COMPILER(MSVC)
+  signal(SIGABRT, &cleanupRoutine);
+#else // ^^^ _CCCL_COMPILER(MSVC) ^^^ / vvv !_CCCL_COMPILER(MSVC)
+  struct sigaction sigabrt_action{};
+  memset(&sigabrt_action, 0, sizeof(sigabrt_action));
+  sigabrt_action.sa_handler = &cleanupRoutine;
+
+  if (sigaction(SIGABRT, &sigabrt_action, nullptr) != 0)
+  {
+    perror("sigaction SIGABRT");
+    exit(EXIT_FAILURE);
+  }
+#endif // !_CCCL_COMPILER(MSVC)
+
+  stream_ctx ctx;
+
+  const int n = 12;
+  double X[n];
+
+  for (int ind = 0; ind < n; ind++)
+  {
+    X[ind] = 1.0 * ind;
+  }
+
+  // This creates a handle that is implicitly a vector of size n
+  auto lX = ctx.logical_data(X);
+
+  auto t = ctx.task(lX.rw());
+  t.start();
+  should_abort = true;
+  t.start();
+  t.end();
+
+  assert(0 && "This should not be reached");
+  return EXIT_FAILURE;
+}

cccl_upstream/cudax/test/stf/error_checks/non_managed_data.cu

@@ -0,0 +1,70 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief Ensure an error is detected when trying to declare a logical data
+ *        with a managed memory data place while the data is not in managed
+ *        memory
+ */
+
+#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+#include <csignal>
+
+using namespace cuda::experimental::stf;
+
+bool should_abort = false;
+
+void cleanupRoutine(int /*unused*/)
+{
+  if (should_abort)
+  {
+    exit(EXIT_SUCCESS);
+  }
+  else
+  {
+    fprintf(stderr, "Unexpected SIGABRT !\n");
+    exit(EXIT_FAILURE);
+  }
+}
+
+int main()
+{
+  /* Setup an handler to catch the SIGABRT signal during the programming error */
+#ifndef NDEBUG
+#  if _CCCL_COMPILER(MSVC)
+  signal(SIGABRT, &cleanupRoutine);
+#  else // ^^^ _CCCL_COMPILER(MSVC) ^^^ / vvv !_CCCL_COMPILER(MSVC)
+  struct sigaction sigabrt_action{};
+  memset(&sigabrt_action, 0, sizeof(sigabrt_action));
+  sigabrt_action.sa_handler = &cleanupRoutine;
+
+  if (sigaction(SIGABRT, &sigabrt_action, nullptr) != 0)
+  {
+    perror("sigaction SIGABRT");
+    exit(EXIT_FAILURE);
+  }
+#  endif // !_CCCL_COMPILER(MSVC)
+
+  stream_ctx ctx;
+
+  logical_data<slice<int>> lX;
+
+  should_abort = true;
+
+  int X[128];
+  lX = ctx.logical_data(X, data_place::managed());
+
+  assert(0 && "This should not be reached");
+  return EXIT_FAILURE;
+#endif
+}

cccl_upstream/cudax/test/stf/error_checks/slice_check_bounds.cu

@@ -0,0 +1,83 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief Ensure that out of bound accesses on slices are detected with the
+ *        CUDASTF_BOUNDSCHECK option set
+ */
+
+/*
+ * We are forcing this option by defining this value to be set.
+ */
+#ifndef NDEBUG
+#  ifndef CUDASTF_BOUNDSCHECK
+#    define CUDASTF_BOUNDSCHECK
+#  endif // CUDASTF_BOUNDSCHECK
+#endif // NDEBUG
+
+#include <cuda/experimental/stf.cuh>
+
+#include <csignal>
+
+using namespace cuda::experimental::stf;
+
+bool should_abort = false;
+
+void cleanupRoutine(int /*unused*/)
+{
+  if (should_abort)
+  {
+    exit(EXIT_SUCCESS);
+  }
+  else
+  {
+    fprintf(stderr, "Unexpected SIGABRT !\n");
+    exit(EXIT_FAILURE);
+  }
+}
+
+int main()
+{
+  /* Setup an handler to catch the SIGABRT signal during the programming error */
+#ifndef NDEBUG
+#  if _CCCL_COMPILER(MSVC)
+  signal(SIGABRT, &cleanupRoutine);
+#  else // ^^^ _CCCL_COMPILER(MSVC) ^^^ / vvv !_CCCL_COMPILER(MSVC)
+  struct sigaction sigabrt_action{};
+  memset(&sigabrt_action, 0, sizeof(sigabrt_action));
+  sigabrt_action.sa_handler = &cleanupRoutine;
+
+  if (sigaction(SIGABRT, &sigabrt_action, nullptr) != 0)
+  {
+    perror("sigaction SIGABRT");
+    exit(EXIT_FAILURE);
+  }
+#  endif // !_CCCL_COMPILER(MSVC)
+
+  context ctx;
+
+  int X[128];
+  logical_data<slice<int>> lX;
+  lX = ctx.logical_data(X);
+
+  should_abort = true;
+
+  // The last access will be out of bounds
+  ctx.parallel_for(lX.shape(), lX.rw())->*[] __device__(size_t i, auto X) {
+    X(i + 1) = 42;
+  };
+
+  ctx.finalize();
+
+  assert(0 && "This should not be reached");
+  return EXIT_FAILURE;
+#endif
+}

cccl_upstream/cudax/test/stf/error_checks/stackable_fence_in_nested.cu

@@ -0,0 +1,75 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2025 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ *
+ * @brief Ensure fence() in a nested stackable context triggers an abort
+ */
+
+#include <cuda/experimental/stf.cuh>
+
+#include <csignal>
+
+using namespace cuda::experimental::stf;
+
+bool should_abort = false;
+
+void cleanupRoutine(int /*unused*/)
+{
+  if (should_abort)
+  {
+    exit(EXIT_SUCCESS);
+  }
+  else
+  {
+    fprintf(stderr, "Unexpected SIGABRT !\n");
+    exit(EXIT_FAILURE);
+  }
+}
+
+int main()
+{
+#if _CCCL_COMPILER(MSVC)
+  signal(SIGABRT, &cleanupRoutine);
+#else // ^^^ _CCCL_COMPILER(MSVC) ^^^ / vvv !_CCCL_COMPILER(MSVC)
+  struct sigaction sigabrt_action{};
+  memset(&sigabrt_action, 0, sizeof(sigabrt_action));
+  sigabrt_action.sa_handler = &cleanupRoutine;
+
+  if (sigaction(SIGABRT, &sigabrt_action, nullptr) != 0)
+  {
+    perror("sigaction SIGABRT");
+    exit(EXIT_FAILURE);
+  }
+#endif // !_CCCL_COMPILER(MSVC)
+
+  stackable_ctx sctx;
+
+  auto lA = sctx.logical_data(shape_of<slice<int>>(64));
+
+  sctx.parallel_for(lA.shape(), lA.write())->*[] __device__(size_t i, auto a) {
+    a(i) = static_cast<int>(i);
+  };
+
+  {
+    auto scope = sctx.graph_scope();
+
+    sctx.parallel_for(lA.shape(), lA.rw())->*[] __device__(size_t i, auto a) {
+      a(i) *= 2;
+    };
+
+    should_abort = true;
+    sctx.fence(); // fence() in nested context must abort
+  }
+
+  _CCCL_ASSERT(false, "This should not be reached");
+  return EXIT_FAILURE;
+}

cccl_upstream/cudax/test/stf/error_checks/stackable_illegal_export.cu

@@ -0,0 +1,82 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ *
+ * @brief Ensure we detect erroneous access on non exportable stackable logical
+ * data after the context was popped
+ *
+ */
+
+#include <cuda/experimental/stf.cuh>
+
+#include <csignal>
+
+using namespace cuda::experimental::stf;
+
+bool should_abort = false;
+
+void cleanupRoutine(int /*unused*/)
+{
+  if (should_abort)
+  {
+    exit(EXIT_SUCCESS);
+  }
+  else
+  {
+    fprintf(stderr, "Unexpected SIGABRT !\n");
+    exit(EXIT_FAILURE);
+  }
+}
+
+int main()
+{
+  /* Setup an handler to catch the SIGABRT signal during the programming error */
+#if _CCCL_COMPILER(MSVC)
+  signal(SIGABRT, &cleanupRoutine);
+#else // ^^^ _CCCL_COMPILER(MSVC) ^^^ / vvv !_CCCL_COMPILER(MSVC)
+  struct sigaction sigabrt_action{};
+  memset(&sigabrt_action, 0, sizeof(sigabrt_action));
+  sigabrt_action.sa_handler = &cleanupRoutine;
+
+  if (sigaction(SIGABRT, &sigabrt_action, nullptr) != 0)
+  {
+    perror("sigaction SIGABRT");
+    exit(EXIT_FAILURE);
+  }
+#endif // !_CCCL_COMPILER(MSVC)
+
+  stackable_ctx sctx;
+  sctx.push();
+
+  auto lB = sctx.logical_data_no_export(shape_of<slice<int>>(1024));
+  lB.set_symbol("B");
+
+  sctx.parallel_for(lB.shape(), lB.write())->*[] __device__(size_t i, auto b) {
+    b(i) = 42;
+  };
+
+  sctx.pop();
+
+  // We are going to try to access B while it was not exportable, and that the
+  // context where it was created has been popped: this should raise an error.
+  should_abort = true;
+
+  sctx.host_launch(lB.read())->*[](auto b) {
+    for (size_t i = 0; i < b.size(); i++)
+    {
+      EXPECT(b(i) == 42);
+    }
+  };
+
+  _CCCL_ASSERT(false, "This should not be reached");
+  return EXIT_FAILURE;
+}

cccl_upstream/cudax/test/stf/error_checks/stackable_inconsistent_access_modes.cu

@@ -0,0 +1,92 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ *
+ * @brief Test that ensures we catch programming errors with inconsistent access modes in nested contexts
+ *
+ * This test verifies that attempting to escalate from read-only to read-write access mode
+ * in nested stackable contexts is properly caught and produces a clear error message.
+ *
+ */
+
+#include <cuda/experimental/stf.cuh>
+
+#include <csignal>
+
+using namespace cuda::experimental::stf;
+
+bool should_abort = false;
+
+void cleanupRoutine(int /*unused*/)
+{
+  if (should_abort)
+  {
+    exit(EXIT_SUCCESS);
+  }
+  else
+  {
+    fprintf(stderr, "Unexpected SIGABRT !\n");
+    exit(EXIT_FAILURE);
+  }
+}
+
+int main()
+{
+  /* Setup an handler to catch the SIGABRT signal during the programming error */
+#if _CCCL_COMPILER(MSVC)
+  signal(SIGABRT, &cleanupRoutine);
+#else // ^^^ _CCCL_COMPILER(MSVC) ^^^ / vvv !_CCCL_COMPILER(MSVC)
+  struct sigaction sigabrt_action{};
+  memset(&sigabrt_action, 0, sizeof(sigabrt_action));
+  sigabrt_action.sa_handler = &cleanupRoutine;
+
+  if (sigaction(SIGABRT, &sigabrt_action, nullptr) != 0)
+  {
+    perror("sigaction SIGABRT");
+    exit(EXIT_FAILURE);
+  }
+#endif // !_CCCL_COMPILER(MSVC)
+
+  stackable_ctx sctx;
+
+  size_t sz = 1024;
+  ::std::vector<int> data(sz);
+
+  // Initialize data
+  for (size_t i = 0; i < sz; i++)
+  {
+    data[i] = static_cast<int>(i);
+  }
+
+  // Create logical data
+  auto ldata = sctx.logical_data(make_slice(data.data(), sz));
+
+  // First scope: push with READ access mode
+  {
+    stackable_ctx::graph_scope_guard scope1{sctx};
+    ldata.push(access_mode::read);
+
+    // We are going to try to escalate from read to rw access mode in nested context:
+    // this should raise an error.
+    should_abort = true;
+
+    // NESTED second scope: attempt to push with RW access mode
+    // This should be caught as an invalid access mode escalation
+    {
+      stackable_ctx::graph_scope_guard scope2{sctx};
+      ldata.push(access_mode::rw); // This should trigger abort()!
+    }
+  }
+
+  _CCCL_ASSERT(false, "This should not be reached");
+  return EXIT_FAILURE;
+}

cccl_upstream/cudax/test/stf/error_checks/stackable_wait_in_nested.cu

@@ -0,0 +1,75 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2025 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ *
+ * @brief Ensure wait() in a nested stackable context triggers an abort
+ */
+
+#include <cuda/experimental/stf.cuh>
+
+#include <csignal>
+
+using namespace cuda::experimental::stf;
+
+bool should_abort = false;
+
+void cleanupRoutine(int /*unused*/)
+{
+  if (should_abort)
+  {
+    exit(EXIT_SUCCESS);
+  }
+  else
+  {
+    fprintf(stderr, "Unexpected SIGABRT !\n");
+    exit(EXIT_FAILURE);
+  }
+}
+
+int main()
+{
+#if _CCCL_COMPILER(MSVC)
+  signal(SIGABRT, &cleanupRoutine);
+#else // ^^^ _CCCL_COMPILER(MSVC) ^^^ / vvv !_CCCL_COMPILER(MSVC)
+  struct sigaction sigabrt_action{};
+  memset(&sigabrt_action, 0, sizeof(sigabrt_action));
+  sigabrt_action.sa_handler = &cleanupRoutine;
+
+  if (sigaction(SIGABRT, &sigabrt_action, nullptr) != 0)
+  {
+    perror("sigaction SIGABRT");
+    exit(EXIT_FAILURE);
+  }
+#endif // !_CCCL_COMPILER(MSVC)
+
+  stackable_ctx sctx;
+
+  auto lval = sctx.logical_data(shape_of<scalar_view<int>>());
+
+  sctx.parallel_for(box(1), lval.write())->*[] __device__(size_t, auto val) {
+    *val = 42;
+  };
+
+  {
+    auto scope = sctx.graph_scope();
+
+    sctx.parallel_for(box(1), lval.rw())->*[] __device__(size_t, auto val) {
+      *val += 1;
+    };
+
+    should_abort = true;
+    sctx.wait(lval); // wait() in nested context must abort
+  }
+
+  _CCCL_ASSERT(false, "This should not be reached");
+  return EXIT_FAILURE;
+}

cccl_upstream/cudax/test/stf/error_checks/uninitialized_data.cu

@@ -0,0 +1,72 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief Ensure an error is detected if we use an uninitialized logical data in a task
+ */
+
+#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+#include <csignal>
+
+using namespace cuda::experimental::stf;
+
+bool should_abort = false;
+
+void cleanupRoutine(int /*unused*/)
+{
+  if (should_abort)
+  {
+    exit(EXIT_SUCCESS);
+  }
+  else
+  {
+    fprintf(stderr, "Unexpected SIGABRT !\n");
+    exit(EXIT_FAILURE);
+  }
+}
+
+int main()
+{
+  /* Setup an handler to catch the SIGABRT signal during the programming error */
+#ifndef NDEBUG
+#  if _CCCL_COMPILER(MSVC)
+  signal(SIGABRT, &cleanupRoutine);
+#  else // ^^^ _CCCL_COMPILER(MSVC) ^^^ / vvv !_CCCL_COMPILER(MSVC)
+  struct sigaction sigabrt_action{};
+  memset(&sigabrt_action, 0, sizeof(sigabrt_action));
+  sigabrt_action.sa_handler = &cleanupRoutine;
+
+  if (sigaction(SIGABRT, &sigabrt_action, nullptr) != 0)
+  {
+    perror("sigaction SIGABRT");
+    exit(EXIT_FAILURE);
+  }
+#  endif // !_CCCL_COMPILER(MSVC)
+
+  stream_ctx ctx;
+
+  logical_data<slice<int>> lX;
+  logical_data<slice<int>> lY;
+
+  int X[128];
+  lX = ctx.logical_data(X);
+
+  should_abort = true;
+
+  // We did not initialize lY, so this task should not be able to use it.
+  ctx.task(lX.rw(), lY.rw())->*[](cudaStream_t, auto, auto) {};
+
+  assert(0 && "This should not be reached");
+  return EXIT_FAILURE;
+#endif
+}

cccl_upstream/cudax/test/stf/error_checks/unsatisfiable_spec.cu

@@ -0,0 +1,72 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief Ensure an error is raised if we try to ask for an unreasonnable
+ *        amount of resources in a thread hierarchy spec
+ */
+
+#include <cuda/experimental/stf.cuh>
+
+#include <csignal>
+
+using namespace cuda::experimental::stf;
+
+bool should_abort = false;
+
+void cleanupRoutine(int /*unused*/)
+{
+  if (should_abort)
+  {
+    exit(EXIT_SUCCESS);
+  }
+  else
+  {
+    fprintf(stderr, "Unexpected SIGABRT !\n");
+    exit(EXIT_FAILURE);
+  }
+}
+
+int main()
+{
+  /* Setup an handler to catch the SIGABRT signal during the programming error */
+#ifndef NDEBUG
+#  if _CCCL_COMPILER(MSVC)
+  signal(SIGABRT, &cleanupRoutine);
+#  else // ^^^ _CCCL_COMPILER(MSVC) ^^^ / vvv !_CCCL_COMPILER(MSVC)
+  struct sigaction sigabrt_action{};
+  memset(&sigabrt_action, 0, sizeof(sigabrt_action));
+  sigabrt_action.sa_handler = &cleanupRoutine;
+
+  if (sigaction(SIGABRT, &sigabrt_action, nullptr) != 0)
+  {
+    perror("sigaction SIGABRT");
+    exit(EXIT_FAILURE);
+  }
+#  endif // !_CCCL_COMPILER(MSVC)
+
+  context ctx;
+
+  int X[128];
+  auto lX = ctx.logical_data(X);
+
+  should_abort = true;
+
+  // We are asking an unreasonnable amount of threads per block
+  auto spec = con(con<128000>());
+  ctx.launch(spec, lX.rw())->*[] __device__(auto th, auto X) {
+    X[th.rank()] = th.rank();
+  };
+
+  assert(0 && "This should not be reached");
+  return EXIT_FAILURE;
+#endif
+}

cccl_upstream/cudax/test/stf/error_checks/write_frozen.cu

@@ -0,0 +1,79 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ *
+ * @brief Ensure temporary data are destroyed
+ *
+ */
+
+#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+#include <csignal>
+
+using namespace cuda::experimental::stf;
+
+bool should_abort = false;
+
+void cleanupRoutine(int /*unused*/)
+{
+  if (should_abort)
+  {
+    exit(EXIT_SUCCESS);
+  }
+  else
+  {
+    fprintf(stderr, "Unexpected SIGABRT !\n");
+    exit(EXIT_FAILURE);
+  }
+}
+
+int main()
+{
+  /* Setup an handler to catch the SIGABRT signal during the programming error */
+#if _CCCL_COMPILER(MSVC)
+  signal(SIGABRT, &cleanupRoutine);
+#else // ^^^ _CCCL_COMPILER(MSVC) ^^^ / vvv !_CCCL_COMPILER(MSVC)
+  struct sigaction sigabrt_action{};
+  memset(&sigabrt_action, 0, sizeof(sigabrt_action));
+  sigabrt_action.sa_handler = &cleanupRoutine;
+
+  if (sigaction(SIGABRT, &sigabrt_action, nullptr) != 0)
+  {
+    perror("sigaction SIGABRT");
+    exit(EXIT_FAILURE);
+  }
+#endif // !_CCCL_COMPILER(MSVC)
+
+  stream_ctx ctx;
+  const int N = 16;
+  int X[N];
+
+  for (int i = 0; i < N; i++)
+  {
+    X[i] = i;
+  }
+
+  auto lX = ctx.logical_data(X);
+
+  lX.freeze(access_mode::rw, data_place::current_device());
+
+  // This is an illegal access because we cannot make a write access on a frozen data
+  should_abort = true;
+
+  ctx.task(lX.rw())->*[](cudaStream_t, auto) {};
+
+  ctx.finalize();
+
+  assert(0 && "This should not be reached");
+  return EXIT_FAILURE;
+}

cccl_upstream/cudax/test/stf/examples/01-axpy-launch-ranges-cg.cu

@@ -0,0 +1,75 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/experimental/__places/partitions/blocked_partition.cuh>
+#include <cuda/experimental/__places/partitions/cyclic_shape.cuh>
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+double X0(int i)
+{
+  return sin((double) i);
+}
+
+double Y0(int i)
+{
+  return cos((double) i);
+}
+
+int main()
+{
+  stream_ctx ctx;
+
+  const int N = 128;
+  double X[N], Y[N];
+
+  for (int ind = 0; ind < N; ind++)
+  {
+    X[ind] = X0(ind);
+    Y[ind] = Y0(ind);
+  }
+
+  const double alpha = 3.14;
+
+  auto handle_X = ctx.logical_data(X, {N});
+  auto handle_Y = ctx.logical_data(Y, {N});
+
+  auto number_devices = 4;
+  auto all_devs       = exec_place::repeat(exec_place::device(0), number_devices);
+
+  auto spec = par(16 * 4, par(4));
+  ctx.launch(spec, all_devs, handle_X.read(), handle_Y.rw())->*[=] _CCCL_DEVICE(auto th, auto x, auto y) {
+    // Blocked partition among elements in the outer most level
+    auto outer_sh = blocked_partition::apply(shape(x), pos4(th.rank(0)), dim4(th.size(0)));
+
+    // Cyclic partition among elements in the remaining levels
+    auto inner_sh = cyclic_partition::apply(outer_sh, pos4(th.inner().rank()), dim4(th.inner().size()));
+
+    for (auto ind : inner_sh)
+    {
+      y(ind) += alpha * x(ind);
+    }
+  };
+
+  ctx.host_launch(handle_X.read(), handle_Y.read())->*[=](auto X, auto Y) {
+    for (int ind = 0; ind < N; ind++)
+    {
+      // Y should be Y0 + alpha X0
+      // fprintf(stderr, "Y[%ld] = %lf - expect %lf\n", ind, Y(ind), (Y0(ind) + alpha * X0(ind)));
+      EXPECT(fabs(Y(ind) - (Y0(ind) + alpha * X0(ind))) < 0.0001);
+
+      // X should be X0
+      EXPECT(fabs(X(ind) - X0(ind)) < 0.0001);
+    }
+  };
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/examples/01-axpy-places.cu

@@ -0,0 +1,119 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ *
+ * @brief This example illustrates how to use the task construct with grids of
+ *        places and composite data places
+ */
+
+#include <cuda/experimental/__places/partitions/tiled_partition.cuh>
+#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+template <typename T>
+__global__ void axpy(size_t start, size_t cnt, T a, const T* x, T* y)
+{
+  int tid      = blockIdx.x * blockDim.x + threadIdx.x;
+  int nthreads = gridDim.x * blockDim.x;
+
+  for (int ind = tid; ind < cnt; ind += nthreads)
+  {
+    y[ind + start] += a * x[ind + start];
+  }
+}
+
+double X0(size_t i)
+{
+  return sin((double) i);
+}
+
+double Y0(size_t i)
+{
+  return cos((double) i);
+}
+
+template <typename Ctx>
+void run()
+{
+  Ctx ctx;
+
+  const int N = 1024 * 1024 * 32;
+  double *X, *Y;
+
+  X = new double[N];
+  Y = new double[N];
+  SCOPE(exit)
+  {
+    delete[] X;
+    delete[] Y;
+  };
+
+  for (size_t ind = 0; ind < N; ind++)
+  {
+    X[ind] = X0(ind);
+    Y[ind] = Y0(ind);
+  }
+
+  //    std::shared_ptr<execution_grid> all_devs = exec_place::all_devices();
+  // use grid [ 0 0 0 0 ] for debugging purpose
+  auto all_devs = exec_place::repeat(exec_place::device(0), 4);
+
+  // 512k doubles = 4MB (2 pages)
+  // A 1D blocking strategy over all devices with a block size of 32 and a round robin distribution of blocks across
+  // devices
+  //    data_place cdp = data_place(exec_place::all_devices().as_grid().get_grid(),
+  //            [](dim4 grid_dim, pos4 index_pos) { return pos4((index_pos.x / (512 * 1024ULL)) % grid_dim.x); });
+
+  data_place cdp = data_place::composite(tiled_partition<512 * 1024ULL>(), all_devs);
+
+  auto handle_X = ctx.logical_data(X, {N});
+  auto handle_Y = ctx.logical_data(Y, {N});
+
+  double alpha = 3.14;
+
+  /* Compute Y = Y + alpha X */
+  auto t = ctx.task(all_devs, handle_X.read(cdp), handle_Y.rw(cdp));
+  t->*[&](auto, auto sX, auto sY) {
+    size_t grid_size = t.grid_dims().size();
+
+    assert(N % grid_size == 0);
+
+    for (size_t i = 0; i < grid_size; i++)
+    {
+      auto active = t.activate_place(i);
+      axpy<<<16, 128, 0, t.get_stream(i)>>>(i * N / grid_size, N / grid_size, alpha, sX.data_handle(), sY.data_handle());
+    }
+  };
+
+  /* Check the result on the host */
+  ctx.host_launch(handle_X.read(), handle_Y.read())->*[&](auto sX, auto sY) {
+    for (size_t ind = 0; ind < N; ind++)
+    {
+      // Y should be Y0 + alpha X0
+      EXPECT(fabs(sY(ind) - (Y0(ind) + alpha * X0(ind))) < 0.0001);
+
+      // X should be X0
+      EXPECT(fabs(sX(ind) - X0(ind)) < 0.0001);
+    }
+  };
+
+  ctx.finalize();
+}
+
+int main()
+{
+  run<stream_ctx>();
+  // Disabled until composite data places are implemented with graphs
+  // run<graph_ctx>();
+}

cccl_upstream/cudax/test/stf/examples/05-stencil-no-copy.cu

@@ -0,0 +1,256 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+/*
+ * DATA BLOCKS
+ *   | GHOSTS | DATA | GHOSTS |
+ */
+template <typename T>
+class data_block
+{
+public:
+  data_block(stream_ctx& ctx, size_t beg, size_t end, size_t GHOST_SIZE)
+      : beg(beg)
+      , end(end)
+      , block_size(end - beg)
+      , ghost_size(GHOST_SIZE)
+      , array(std::vector<T>(block_size + 2 * ghost_size))
+      , handle(ctx.logical_data(&array[0], block_size + 2 * ghost_size))
+  {}
+
+public:
+  size_t beg;
+  size_t end;
+  size_t block_size;
+  size_t ghost_size;
+  int dev_id;
+
+private:
+  std::vector<T> array;
+
+public:
+  // HANDLE = whole data + boundaries
+  logical_data<slice<T>> handle;
+};
+
+template <typename T>
+T check_sum(stream_ctx& ctx, data_block<T>& bn)
+{
+  T sum = 0.0;
+
+  auto t = ctx.task(exec_place::host(), bn.handle.read());
+  t->*[&](cudaStream_t stream, auto h_center) {
+    cuda_safe_call(cudaStreamSynchronize(stream));
+    for (size_t offset = bn.ghost_size; offset < bn.ghost_size + bn.block_size; offset++)
+    {
+      sum += h_center.data_handle()[offset];
+    }
+  };
+
+  return sum;
+}
+
+// array and array1 have a size of (cnt + 2*ghost_size)
+template <typename T>
+__global__ void stencil_kernel(size_t cnt, size_t ghost_size, T* array, const T* array1)
+{
+  for (size_t idx = threadIdx.x + blockIdx.x * blockDim.x; idx < cnt; idx += blockDim.x * gridDim.x)
+  {
+    size_t idx2 = idx + ghost_size;
+    array[idx2] = 0.9 * array1[idx2] + 0.05 * array1[idx2 - 1] + 0.05 * array1[idx2 + 1];
+  }
+}
+
+template <typename T>
+void stencil(stream_ctx& ctx, data_block<T>& bn, data_block<T>& bn1)
+{
+  int dev = bn.dev_id;
+
+  auto t = ctx.task(exec_place::device(dev), bn.handle.rw(), bn1.handle.read());
+  t->*[&](cudaStream_t stream, auto bn_array, auto bn1_array) {
+    stencil_kernel<T>
+      <<<32, 64, 0, stream>>>(bn.block_size, bn.ghost_size, bn_array.data_handle(), bn1_array.data_handle());
+  };
+}
+
+template <typename T>
+__global__ void copy_kernel(size_t cnt, T* dst, const T* src)
+{
+  for (size_t idx = threadIdx.x + blockIdx.x * blockDim.x; idx < cnt; idx += blockDim.x * gridDim.x)
+  {
+    dst[idx] = src[idx];
+  }
+}
+
+template <typename T>
+void copy_task(
+  stream_ctx& ctx,
+  size_t cnt,
+  logical_data<slice<T>>& dst,
+  size_t offset_dst,
+  int dst_dev,
+  logical_data<slice<T>>& src,
+  size_t offset_src,
+  int src_dev)
+{
+  auto t = ctx.task(exec_place::device(dst_dev), dst.rw(), src.read(data_place::device(src_dev)));
+  t->*[&](cudaStream_t stream, auto dst_array, auto src_array) {
+    int nblocks = (cnt > 64) ? 32 : 1;
+    copy_kernel<T>
+      <<<nblocks, 64, 0, stream>>>(cnt, dst_array.data_handle() + offset_dst, src_array.data_handle() + offset_src);
+  };
+}
+
+// Copy left/right handles from neighbours to the array
+template <typename T>
+void update_halo(stream_ctx& ctx, data_block<T>& bn, data_block<T>& left, data_block<T>& right)
+{
+  size_t gs = bn.ghost_size;
+  size_t bs = bn.block_size;
+
+  // Copy the bn.ghost_size last computed items in "left" (outside the halo)
+  copy_task<T>(ctx, gs, bn.handle, 0, bn.dev_id, left.handle, bs, left.dev_id);
+
+  // Copy the bn.ghost_size first computed items (outside the halo)
+  copy_task<T>(ctx, gs, bn.handle, gs + bs, bn.dev_id, right.handle, gs, right.dev_id);
+}
+
+// Copy inner part of bn into bn1
+template <typename T>
+void copy_inner(stream_ctx& ctx, data_block<T>& bn1, data_block<T>& bn)
+{
+  size_t gs = bn.ghost_size;
+  size_t bs = bn.block_size;
+
+  int dev_id = bn.dev_id;
+
+  // Copy the bn.ghost_size last computed items in "left" (outside the halo)
+  copy_task<T>(ctx, bs, bn1.handle, gs, dev_id, bn.handle, gs, dev_id);
+}
+
+int main(int argc, char** argv)
+{
+  int ndevs;
+  cuda_safe_call(cudaGetDeviceCount(&ndevs));
+
+  stream_ctx ctx;
+
+  int NITER         = 500;
+  size_t NBLOCKS    = 4 * ndevs;
+  size_t BLOCK_SIZE = 2048 * 1024;
+
+  if (argc > 1)
+  {
+    NITER = atoi(argv[1]);
+  }
+
+  if (argc > 2)
+  {
+    NBLOCKS = atoi(argv[2]);
+  }
+
+  const size_t GHOST_SIZE = 1;
+
+  size_t TOTAL_SIZE = NBLOCKS * BLOCK_SIZE;
+
+  double* U0 = new double[NBLOCKS * BLOCK_SIZE];
+  for (size_t idx = 0; idx < NBLOCKS * BLOCK_SIZE; idx++)
+  {
+    U0[idx] = (idx == 0) ? 1.0 : 0.0;
+  }
+
+  std::vector<data_block<double>> Un;
+  std::vector<data_block<double>> Un1;
+
+  // Create blocks and allocates host data
+  for (size_t b = 0; b < NBLOCKS; b++)
+  {
+    size_t beg = b * BLOCK_SIZE;
+    size_t end = (b + 1) * BLOCK_SIZE;
+
+    Un.emplace_back(ctx, beg, end, 1ull);
+    Un1.emplace_back(ctx, beg, end, 1ull);
+  }
+
+  for (size_t b = 0; b < NBLOCKS; b++)
+  {
+    Un[b].dev_id  = b % ndevs;
+    Un1[b].dev_id = b % ndevs;
+  }
+
+  // Fill blocks with initial values. For the sake of simplicity, we are
+  // using a synchronization primitive and host code, but this could have
+  // been written asynchronously using host callbacks.
+  for (size_t b = 0; b < NBLOCKS; b++)
+  {
+    size_t beg = b * BLOCK_SIZE;
+
+    auto t = ctx.task(exec_place::host(), Un[b].handle.rw(), Un1[b].handle.rw());
+    t->*[&](cudaStream_t stream, auto Un_vals, auto Un1_vals) {
+      cuda_safe_call(cudaStreamSynchronize(stream));
+      for (size_t local_idx = 0; local_idx < BLOCK_SIZE; local_idx++)
+      {
+        double val                                     = U0[(beg + local_idx + TOTAL_SIZE) % TOTAL_SIZE];
+        Un1_vals.data_handle()[local_idx + GHOST_SIZE] = val;
+        Un_vals.data_handle()[local_idx + GHOST_SIZE]  = val;
+      }
+    };
+  }
+
+  for (int iter = 0; iter < NITER; iter++)
+  {
+    for (size_t b = 0; b < NBLOCKS; b++)
+    {
+      update_halo(ctx, Un1[b], Un[(b - 1 + NBLOCKS) % NBLOCKS], Un[(b + 1) % NBLOCKS]);
+    }
+
+    // UPDATE Un from Un1
+    for (size_t b = 0; b < NBLOCKS; b++)
+    {
+      stencil(ctx, Un[b], Un1[b]);
+    }
+
+#if 0
+    // We make sure that the total sum of elements remains constant
+    if (iter % 250 == 0)
+    {
+      double sum = 0.0;
+      for (size_t b = 0; b < NBLOCKS; b++)
+      {
+        sum += check_sum(ctx, Un[b]);
+      }
+
+      // fprintf(stderr, "iter %d : CHECK SUM = %e\n", iter, sum);
+    }
+#endif
+
+    for (size_t b = 0; b < NBLOCKS; b++)
+    {
+      // Copy inner part of Un into Un1
+      copy_inner(ctx, Un[b], Un1[b]);
+    }
+  }
+
+  // In this stencil, the sum of the elements is supposed to be a constant
+  double sum = 0.0;
+  for (size_t b = 0; b < NBLOCKS; b++)
+  {
+    sum += check_sum(ctx, Un[b]);
+  }
+
+  double err = fabs(sum - 1.0);
+  EXPECT(err < 0.0001);
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/examples/05-stencil-places.cu

@@ -0,0 +1,92 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/experimental/__places/partitions/tiled_partition.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+template <typename T>
+__global__ void stencil_kernel(slice<T> Un, slice<const T> Un1)
+{
+  size_t N = Un.extent(0);
+  for (size_t i = threadIdx.x + blockIdx.x * blockDim.x; i < N; i += blockDim.x * gridDim.x)
+  {
+    Un(i) = 0.9 * Un1(i) + 0.05 * Un1((i + N - 1) % N) + 0.05 * Un1((i + 1) % N);
+  }
+}
+
+int main(int argc, char** argv)
+{
+  stream_ctx ctx;
+
+  int NITER               = 500;
+  int NBLOCKS             = 20;
+  const size_t BLOCK_SIZE = 2048 * 1024;
+
+  if (argc > 1)
+  {
+    NITER = atoi(argv[1]);
+  }
+
+  if (argc > 2)
+  {
+    NBLOCKS = atoi(argv[2]);
+  }
+
+  const size_t TOTAL_SIZE = NBLOCKS * BLOCK_SIZE;
+
+  double* Un  = new double[TOTAL_SIZE];
+  double* Un1 = new double[TOTAL_SIZE];
+
+  for (size_t idx = 0; idx < TOTAL_SIZE; idx++)
+  {
+    Un[idx]  = (idx == 0) ? 1.0 : 0.0;
+    Un1[idx] = Un[idx];
+  }
+
+  auto lUn  = ctx.logical_data(make_slice(Un, TOTAL_SIZE));
+  auto lUn1 = ctx.logical_data(make_slice(Un1, TOTAL_SIZE));
+
+  //    std::shared_ptr<execution_grid> all_devs = exec_place::all_devices();
+  // use grid [ 0 0 0 0 ] for debugging purpose
+  auto all_devs = exec_place::repeat(exec_place::device(0), 4);
+
+  data_place cdp = data_place::composite(tiled_partition<BLOCK_SIZE>(), all_devs);
+
+  for (int iter = 0; iter < NITER; iter++)
+  {
+    // UPDATE Un from Un1
+    ctx.task(lUn.rw(cdp), lUn1.read(cdp))->*[&](auto stream, auto sUn, auto sUn1) {
+      stencil_kernel<double><<<32, 128, 0, stream>>>(sUn, sUn1);
+    };
+
+    // We make sure that the total sum of elements remains constant
+    if (iter % 250 == 0)
+    {
+      double sum = 0.0;
+
+      ctx.task(exec_place::host(), lUn.read())->*[&](auto stream, auto sUn) {
+        cuda_safe_call(cudaStreamSynchronize(stream));
+        for (size_t offset = 0; offset < TOTAL_SIZE; offset++)
+        {
+          sum += sUn(offset);
+        }
+      };
+
+      // TODO add an assertion to check whether sum is close enough to 1.0
+      // fprintf(stderr, "iter %d : CHECK SUM = %e\n", iter, sum);
+    }
+
+    std::swap(lUn, lUn1);
+  }
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/examples/05-stencil.cu

@@ -0,0 +1,265 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+static stream_ctx ctx;
+
+/*
+ * DATA BLOCKS
+ *   | GHOSTS | DATA | GHOSTS |
+ */
+template <typename T>
+class data_block
+{
+public:
+  data_block(size_t beg, size_t end, size_t GHOST_SIZE)
+      : beg(beg)
+      , end(end)
+      , block_size(end - beg)
+      , ghost_size(GHOST_SIZE)
+      , array(std::vector<T>(block_size + 2 * ghost_size))
+      , left_interface(std::vector<T>(ghost_size))
+      , right_interface(std::vector<T>(ghost_size))
+      , handle(ctx.logical_data(&array[0], block_size + 2 * ghost_size))
+      , left_handle(ctx.logical_data(&left_interface[0], ghost_size))
+      , right_handle(ctx.logical_data(&right_interface[0], ghost_size))
+  {}
+
+  T check_sum()
+  {
+    T sum = 0.0;
+
+    ctx.task(exec_place::host(), handle.read())->*[&](cudaStream_t stream, auto sn) {
+      cuda_safe_call(cudaStreamSynchronize(stream));
+      const T* h_center = sn.data_handle();
+      for (size_t offset = ghost_size; offset < ghost_size + block_size; offset++)
+      {
+        sum += h_center[offset];
+      }
+    };
+
+    return sum;
+  }
+
+public:
+  size_t beg;
+  size_t end;
+  size_t block_size;
+  size_t ghost_size;
+  int preferred_device;
+
+private:
+  std::vector<T> array;
+  std::vector<T> left_interface;
+  std::vector<T> right_interface;
+
+public:
+  // HANDLE = whole data + boundaries
+  logical_data<slice<T>> handle;
+  // A piece of data to store the left part of the block
+  logical_data<slice<T>> left_handle;
+  // A piece of data to store the right part of the block
+  logical_data<slice<T>> right_handle;
+};
+
+// array and array1 have a size of (cnt + 2*ghost_size)
+template <typename T>
+__global__ void stencil_kernel(size_t cnt, size_t ghost_size, T* array, const T* array1)
+{
+  for (size_t idx = threadIdx.x + blockIdx.x * blockDim.x; idx < cnt; idx += blockDim.x * gridDim.x)
+  {
+    size_t idx2 = idx + ghost_size;
+    array[idx2] = 0.9 * array1[idx2] + 0.05 * array1[idx2 - 1] + 0.05 * array1[idx2 + 1];
+  }
+}
+
+// bn1.array = bn.array
+template <typename T>
+void stencil(data_block<T>& bn, data_block<T>& bn1)
+{
+  int dev = bn.preferred_device;
+
+  ctx.task(exec_place::device(dev), bn.handle.rw(), bn1.handle.read())->*[&](cudaStream_t stream, auto sN, auto sN1) {
+    stencil_kernel<T><<<32, 64, 0, stream>>>(bn.block_size, bn.ghost_size, sN.data_handle(), sN1.data_handle());
+  };
+}
+
+template <typename T>
+__global__ void copy_kernel(size_t cnt, T* dst, const T* src)
+{
+  for (size_t idx = threadIdx.x + blockIdx.x * blockDim.x; idx < cnt; idx += blockDim.x * gridDim.x)
+  {
+    dst[idx] = src[idx];
+  }
+}
+
+template <typename T>
+void copy_task(
+  size_t cnt, logical_data<slice<T>>& dst, size_t offset_dst, logical_data<slice<T>>& src, size_t offset_src, int dev)
+{
+  ctx.task(exec_place::device(dev), dst.rw(), src.read())->*[&](cudaStream_t stream, auto dstS, auto srcS) {
+    int nblocks = (cnt > 64) ? 32 : 1;
+    copy_kernel<T><<<nblocks, 64, 0, stream>>>(cnt, dstS.data_handle() + offset_dst, srcS.data_handle() + offset_src);
+  };
+}
+
+template <typename T>
+void update_inner_interfaces(data_block<T>& bn)
+{
+  // LEFT
+  copy_task<T>(bn.ghost_size, bn.left_handle, 0, bn.handle, bn.ghost_size, bn.preferred_device);
+
+  // RIGHT
+  copy_task<T>(bn.ghost_size, bn.right_handle, 0, bn.handle, bn.block_size, bn.preferred_device);
+}
+
+// Copy left/right handles from neighbours to the array
+template <typename T>
+void update_outer_interfaces(data_block<T>& bn, data_block<T>& left, data_block<T>& right)
+{
+  // update_outer_interface_left
+  copy_task<T>(bn.ghost_size, bn.handle, 0, left.right_handle, 0, bn.preferred_device);
+
+  // update_outer_interface_right
+  copy_task<T>(bn.ghost_size, bn.handle, bn.ghost_size + bn.block_size, right.left_handle, 0, bn.preferred_device);
+}
+
+// bn1.array = bn.array
+template <typename T>
+void copy_array(data_block<T>& bn, data_block<T>& bn1)
+{
+  assert(bn.preferred_device == bn1.preferred_device);
+  copy_task<T>(bn.block_size + 2 * bn.ghost_size, bn1.handle, 0, bn.handle, 0, bn.preferred_device);
+}
+
+int main(int argc, char** argv)
+{
+  int NITER         = 500;
+  size_t NBLOCKS    = 4;
+  size_t BLOCK_SIZE = 2048 * 1024;
+
+  if (argc > 1)
+  {
+    NITER = atoi(argv[1]);
+  }
+
+  if (argc > 2)
+  {
+    NBLOCKS = atoi(argv[2]);
+  }
+
+  const size_t GHOST_SIZE = 1;
+
+  size_t TOTAL_SIZE = NBLOCKS * BLOCK_SIZE;
+
+  int ndevs;
+  cuda_safe_call(cudaGetDeviceCount(&ndevs));
+
+  // fprintf(stderr, "GOT %d devices\n", ndevs);
+
+  double* U0 = new double[NBLOCKS * BLOCK_SIZE];
+  for (size_t idx = 0; idx < NBLOCKS * BLOCK_SIZE; idx++)
+  {
+    U0[idx] = (idx == 0) ? 1.0 : 0.0;
+  }
+
+  std::vector<data_block<double>> Un;
+  std::vector<data_block<double>> Un1;
+
+  // Create blocks and allocates host data
+  for (size_t b = 0; b < NBLOCKS; b++)
+  {
+    size_t beg = b * BLOCK_SIZE;
+    size_t end = (b + 1) * BLOCK_SIZE;
+
+    Un.emplace_back(beg, end, 1ull);
+    Un1.emplace_back(beg, end, 1ull);
+  }
+
+  for (size_t b = 0; b < NBLOCKS; b++)
+  {
+    Un[b].preferred_device  = b % ndevs;
+    Un1[b].preferred_device = b % ndevs;
+  }
+
+  // Fill blocks with initial values. For the sake of simplicity, we are
+  // using a synchronization primitive and host code, but this could have
+  // been written asynchronously using host callbacks.
+  for (size_t b = 0; b < NBLOCKS; b++)
+  {
+    size_t beg = b * BLOCK_SIZE;
+
+    ctx.task(exec_place::host(), Un1[b].handle.rw())->*[&](cudaStream_t stream, auto sUn1) {
+      cuda_safe_call(cudaStreamSynchronize(stream));
+      double* Un1_vals = sUn1.data_handle();
+
+      for (size_t local_idx = 0; local_idx < BLOCK_SIZE; local_idx++)
+      {
+        Un1_vals[local_idx + GHOST_SIZE] = U0[(beg + local_idx + TOTAL_SIZE) % TOTAL_SIZE];
+      }
+    };
+  }
+
+  for (int iter = 0; iter < NITER; iter++)
+  {
+    for (size_t b = 0; b < NBLOCKS; b++)
+    {
+      // Update the internal copies of the left and right boundaries
+      update_inner_interfaces(Un1[b]);
+    }
+
+    for (size_t b = 0; b < NBLOCKS; b++)
+    {
+      // Apply ghost cells from neighbours to put then in the "center" array
+      update_outer_interfaces(Un1[b], Un1[(b - 1 + NBLOCKS) % NBLOCKS], Un1[(b + 1) % NBLOCKS]);
+    }
+
+    // UPDATE Un from Un1
+    for (size_t b = 0; b < NBLOCKS; b++)
+    {
+      stencil(Un[b], Un1[b]);
+    }
+
+    for (size_t b = 0; b < NBLOCKS; b++)
+    {
+      // Save Un into Un1
+      copy_array(Un[b], Un1[b]);
+    }
+
+#if 0
+    // We make sure that the total sum of elements remains constant
+    if (iter % 250 == 0)
+    {
+      double check_sum = 0.0;
+      for (size_t b = 0; b < NBLOCKS; b++)
+      {
+        check_sum += Un[b].check_sum();
+      }
+
+      // fprintf(stderr, "iter %d : CHECK SUM = %e\n", iter, check_sum);
+    }
+#endif
+  }
+
+  // In this stencil, the sum of the elements is supposed to be a constant
+  double check_sum = 0.0;
+  for (size_t b = 0; b < NBLOCKS; b++)
+  {
+    check_sum += Un[b].check_sum();
+  }
+
+  double err = fabs(check_sum - 1.0);
+  EXPECT(err < 0.0001);
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/examples/05-stencil2d-places.cu

@@ -0,0 +1,114 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/experimental/__places/partitions/tiled_partition.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+#include <cuda/experimental/__stf/utility/pretty_print.cuh>
+
+using namespace cuda::experimental::stf;
+
+template <typename T>
+__global__ void stencil2D_kernel(slice<T, 2> sUn, slice<const T, 2> sUn1)
+{
+  size_t N = sUn.extent(0);
+  for (size_t i = threadIdx.x + blockIdx.x * blockDim.x; i < N; i += blockDim.x * gridDim.x)
+  {
+    for (size_t j = 0; j < N; j++)
+    {
+      sUn(j, i) = 0.8 * sUn1(j, i) + 0.05 * sUn1(j, (i + 1) % N) + 0.05 * sUn1(j, (i - 1 + N) % N)
+                + 0.05 * sUn1((j + 1) % N, i) + 0.05 * sUn1((j - 1 + N) % N, i);
+    }
+  }
+}
+
+int main(int argc, char** argv)
+{
+  stream_ctx ctx;
+
+  size_t NITER  = 500;
+  size_t N      = 1000;
+  bool vtk_dump = false;
+
+  if (argc > 1)
+  {
+    NITER = atoi(argv[1]);
+  }
+
+  if (argc > 2)
+  {
+    N = atoi(argv[2]);
+  }
+
+  if (argc > 3)
+  {
+    int val  = atoi(argv[3]);
+    vtk_dump = (val == 1);
+  }
+
+  size_t TOTAL_SIZE = N * N;
+
+  double* Un  = new double[TOTAL_SIZE];
+  double* Un1 = new double[TOTAL_SIZE];
+
+  for (size_t idx = 0; idx < TOTAL_SIZE; idx++)
+  {
+    Un[idx]  = (idx == 0) ? 1.0 : 0.0;
+    Un1[idx] = Un[idx];
+  }
+
+  auto lUn  = ctx.logical_data(make_slice(Un, std::tuple{N, N}, N));
+  auto lUn1 = ctx.logical_data(make_slice(Un1, std::tuple{N, N}, N));
+
+  //    std::shared_ptr<execution_grid> all_devs = exec_place::all_devices();
+  // use grid [ 0 0 0 0 ] for debugging purpose
+  auto all_devs = exec_place::repeat(exec_place::device(0), 4);
+
+  // Partition over the vector of processor along the y-axis of the data domain
+  // TODO implement the proper tiled_partitioning along y !
+  data_place cdp = data_place::composite(tiled_partition<128>(), all_devs);
+
+  for (size_t iter = 0; iter < NITER; iter++)
+  {
+    // UPDATE Un from Un1
+    ctx.task(lUn.rw(cdp), lUn1.read(cdp))->*[&](auto stream, auto sUn, auto sUn1) {
+      stencil2D_kernel<double><<<32, 128, 0, stream>>>(sUn, sUn1);
+    };
+
+    // We make sure that the total sum of elements remains constant
+    if (iter % 250 == 0)
+    {
+      double sum = 0.0;
+
+      ctx.task(exec_place::host(), lUn.read())->*[&](auto stream, auto sUn) {
+        cuda_safe_call(cudaStreamSynchronize(stream));
+        for (size_t j = 0; j < N; j++)
+        {
+          for (size_t i = 0; i < N; i++)
+          {
+            sum += sUn(j, i);
+          }
+        }
+
+        if (vtk_dump)
+        {
+          char str[32];
+          snprintf(str, 32, "Un_%05zu.vtk", iter);
+          mdspan_to_vtk(sUn, std::string(str));
+        }
+      };
+
+      // fprintf(stderr, "iter %d : CHECK SUM = %e\n", iter, sum);
+    }
+
+    std::swap(lUn, lUn1);
+  }
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/examples/07-cholesky-redux.cu

@@ -0,0 +1,750 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ *
+ * @brief This example implements a Cholesky decomposition over multiple devices using CUBLAS and CUSOLVER
+ *
+ * It also illustrates how we can use CUDASTF to allocate temporary data for CUSOLVER in CUDASTF tasks
+ */
+
+#include <cuda/experimental/__stf/stream/interfaces/slice_reduction_ops.cuh>
+#include <cuda/experimental/__stf/utility/nvtx.cuh>
+#include <cuda/experimental/stf.cuh>
+
+#include <iostream>
+
+#define TILED
+
+using namespace cuda::experimental::stf;
+
+// Global for the sake of simplicity !
+stream_ctx ctx;
+
+/* Get a CUBLAS handle valid on the current execution place, or initialize it lazily */
+cublasHandle_t& get_cublas_handle(const exec_place& ep = exec_place::current_device())
+{
+  static std::unordered_map<exec_place, cublasHandle_t, hash<exec_place>> cublas_handles;
+  auto& result = cublas_handles[ep];
+  if (result == cublasHandle_t())
+  { // not found, default value inserted
+    // Lazy initialization, and save the handle for future use
+    cuda_safe_call(cublasCreate(&result));
+  }
+  return result;
+}
+
+/* Get a CUSOLVER handle valid on the current execution place, or initialize it lazily */
+cusolverDnHandle_t& get_cusolver_handle(const exec_place& ep = exec_place::current_device())
+{
+  static std::unordered_map<exec_place, cusolverDnHandle_t, hash<exec_place>> cusolver_handles;
+  auto& result = cusolver_handles[ep];
+  if (result == cusolverDnHandle_t())
+  { // not found, default value inserted
+    // Lazy initialization, and save the handle for future use
+    cuda_safe_call(cusolverDnCreate(&result));
+  }
+  return result;
+}
+
+template <typename T>
+class matrix
+{
+public:
+  matrix(int NROWS, int NCOLS, int BLOCKSIZE_ROWS, int BLOCKSIZE_COLS, bool is_sym, const char* _symbol = "matrix")
+  {
+    symbol = _symbol;
+
+    sym_matrix = is_sym;
+
+    m  = NROWS;
+    mb = BLOCKSIZE_ROWS;
+
+    n  = NCOLS;
+    nb = BLOCKSIZE_COLS;
+
+    assert(m % mb == 0);
+    assert(n % nb == 0);
+
+    // cuda_safe_call(cudaMallocHost(&h_array, m*n*sizeof(T)));
+    // fprintf(stderr, "Allocating %ld x %ld x %ld = %ld bytes (%f GB) on host for %s\n", m, n, sizeof(T), s,
+    //        s / (1024.0 * 1024.0 * 1024.0), _symbol);
+    h_array.resize(m * n);
+    cuda_safe_call(cudaHostRegister(&h_array[0], h_array.size() * sizeof(T), cudaHostRegisterPortable));
+
+    // Compute the number of blocks
+    mt = m / mb;
+    nt = n / nb;
+
+    handles.resize(mt * nt);
+
+    for (size_t colb = 0; colb < nt; colb++)
+    {
+      size_t low_rowb = sym_matrix ? colb : 0;
+      for (size_t rowb = low_rowb; rowb < mt; rowb++)
+      {
+        T* addr_h = get_block_h(rowb, colb);
+        auto& h   = handle(rowb, colb);
+
+#ifdef TILED
+        // tiles are stored contiguously
+        size_t ld = mb;
+#else
+        size_t ld = m;
+#endif
+        std::ignore = ld; // work around bug in compiler
+        h           = ctx.logical_data(make_slice(addr_h, std::tuple{mb, nb}, ld));
+        h.set_symbol(std::string(symbol) + "_" + std::to_string(rowb) + "_" + std::to_string(colb));
+        h.set_write_back(false);
+      }
+    }
+
+    cuda_safe_call(cudaGetDeviceCount(&ndevs));
+    for (int a = 1; a * a <= ndevs; a++)
+    {
+      if (ndevs % a == 0)
+      {
+        grid_p = a;
+        grid_q = ndevs / a;
+      }
+    }
+
+    assert(grid_p * grid_q == ndevs);
+
+    // std::cout << "FOUND " << ndevs << " DEVICES "
+    //          << "p=" << grid_p << " q=" << grid_q << '\n';
+  }
+
+  int get_preferred_devid(int row, int col)
+  {
+    return (row % grid_p) + (col % grid_q) * grid_p;
+  }
+
+  auto& handle(int row, int col)
+  {
+    return handles[row + col * mt];
+  }
+
+  size_t get_index(size_t row, size_t col)
+  {
+#ifdef TILED
+    // Find which tile contains this element
+    int tile_row = row / mb;
+    int tile_col = col / nb;
+
+    size_t tile_size = mb * nb;
+
+    // Look for the index of the beginning of the tile
+    size_t tile_start = (tile_row + mt * tile_col) * tile_size;
+
+    // Offset within the tile
+    size_t offset = (row % mb) + (col % nb) * mb;
+
+    return tile_start + offset;
+#else
+    return row + col * m;
+#endif
+  }
+
+  T* get_block_h(int brow, int bcol)
+  {
+    size_t index = get_index(brow * mb, bcol * nb);
+    return &h_array[index];
+  }
+
+  // Fill with func(Matrix*,row, col)
+  template <typename Fun>
+  void fill(Fun&& fun)
+  {
+    nvtx_range r("fill");
+    // Fill blocks by blocks
+    for (size_t colb = 0; colb < nt; colb++)
+    {
+      size_t low_rowb = sym_matrix ? colb : 0;
+      for (size_t rowb = low_rowb; rowb < mt; rowb++)
+      {
+        // Each task fills a block
+        auto& h   = handle(rowb, colb);
+        int devid = get_preferred_devid(rowb, colb);
+
+        ctx.parallel_for(exec_place::device(devid), h.shape(), h.write()).set_symbol("INIT")->*
+          [=] __device__(size_t lrow, size_t lcol, auto sA) {
+            size_t row     = lrow + rowb * sA.extent(0);
+            size_t col     = lcol + colb * sA.extent(1);
+            sA(lrow, lcol) = fun(row, col);
+          };
+      }
+    }
+  }
+
+  std::vector<T> h_array;
+  size_t m; // nrows
+  size_t n; // ncols
+
+  // Is this a sym matrix ? (lower assumed)
+  bool sym_matrix;
+
+  size_t mb; // block size (rows)
+  size_t nb; // block size (cols)
+
+  size_t mt; // number of column blocks
+  size_t nt; // number of row blocks
+
+  // abstract data handles
+  std::vector<logical_data<slice<double, 2>>> handles;
+
+  const char* symbol;
+
+  // for the mapping
+  int ndevs;
+  int grid_p, grid_q;
+};
+
+void DPOTRF(cublasFillMode_t uplo, class matrix<double>& A, int A_row, int A_col)
+{
+  auto& Akk    = A.handle(A_row, A_col);
+  size_t m_akk = Akk.shape().extent(0);
+  // Note that the handle may be different from the actual handle...
+  int Lwork_expected;
+  cuda_safe_call(cusolverDnDpotrf_bufferSize(get_cusolver_handle(), uplo, m_akk, nullptr, 0, &Lwork_expected));
+
+  auto potrf_buffer = ctx.logical_data<double>(Lwork_expected);
+  auto devInfo      = ctx.logical_data(shape_of<slice<int>>(1));
+
+  auto t =
+    ctx.task(exec_place::device(A.get_preferred_devid(A_row, A_col)), Akk.rw(), potrf_buffer.write(), devInfo.write());
+  t.set_symbol("DPOTRF");
+  t->*[uplo](cudaStream_t s, auto sAkk, auto buffer, auto info) {
+    auto& h = get_cusolver_handle();
+    cuda_safe_call(cusolverDnSetStream(h, s));
+
+    cuda_safe_call(cusolverDnDpotrf(
+      h,
+      uplo,
+      sAkk.extent(0),
+      sAkk.data_handle(),
+      sAkk.stride(1),
+      buffer.data_handle(),
+      buffer.extent(0),
+      info.data_handle()));
+  };
+}
+
+void DGEMM(
+  cublasOperation_t transa,
+  cublasOperation_t transb,
+  double alpha,
+  class matrix<double>& A,
+  int A_row,
+  int A_col,
+  class matrix<double>& B,
+  int B_row,
+  int B_col,
+  double beta,
+  class matrix<double>& C,
+  int C_row,
+  int C_col)
+{
+  auto redux_op = std::make_shared<slice_reduction_op_sum<double, 2>>();
+
+  // If beta == 1.0 (we assume this is exactly 1.0), then this operation is
+  // an accumulation with the add operator
+  auto dep_c = (beta == 1.0) ? C.handle(C_row, C_col).relaxed(redux_op) : C.handle(C_row, C_col).rw();
+  auto t     = ctx.task(exec_place::device(A.get_preferred_devid(C_row, C_col)),
+                    A.handle(A_row, A_col).read(),
+                    B.handle(B_row, B_col).read(),
+                    dep_c);
+  t.set_symbol("DGEMM");
+  t->*[transa, transb, alpha, beta](cudaStream_t s, auto sA, auto sB, auto sC) {
+    EXPECT(sC.data_handle() != nullptr);
+    auto& h = get_cublas_handle();
+    cuda_safe_call(cublasSetStream(h, s));
+
+    auto k = (transa == CUBLAS_OP_N) ? sA.extent(1) : sA.extent(0);
+    cuda_safe_call(cublasDgemm(
+      h,
+      transa,
+      transb,
+      sC.extent(0),
+      sC.extent(1),
+      k,
+      &alpha,
+      sA.data_handle(),
+      sA.stride(1),
+      sB.data_handle(),
+      sB.stride(1),
+      &beta,
+      sC.data_handle(),
+      sC.stride(1)));
+  };
+}
+
+void DSYRK(
+  cublasFillMode_t uplo,
+  cublasOperation_t trans,
+  double alpha,
+  class matrix<double>& A,
+  int A_row,
+  int A_col,
+  double beta,
+  class matrix<double>& C,
+  int C_row,
+  int C_col)
+{
+  auto t = ctx.task(exec_place::device(A.get_preferred_devid(C_row, C_col)),
+                    A.handle(A_row, A_col).read(),
+                    C.handle(C_row, C_col).rw());
+  t.set_symbol("DSYRK");
+  t->*[uplo, trans, alpha, beta](cudaStream_t s, auto sA, auto sC) {
+    auto& h = get_cublas_handle();
+    cuda_safe_call(cublasSetStream(h, s));
+
+    // number of rows of matrix op(A) and C
+    auto n = sC.extent(0);
+
+    // number of columns of matrix op(A)
+    auto k = (trans == CUBLAS_OP_N) ? sA.extent(1) : sA.extent(0);
+
+    cuda_safe_call(
+      cublasDsyrk(h, uplo, trans, n, k, &alpha, sA.data_handle(), sA.stride(1), &beta, sC.data_handle(), sC.stride(1)));
+  };
+}
+
+void DTRSM(
+  cublasSideMode_t side,
+  cublasFillMode_t uplo,
+  cublasOperation_t transa,
+  cublasDiagType_t diag,
+  double alpha,
+  class matrix<double>& A,
+  int A_row,
+  int A_col,
+  class matrix<double>& B,
+  int B_row,
+  int B_col)
+{
+  auto t = ctx.task(exec_place::device(A.get_preferred_devid(B_row, B_col)),
+                    A.handle(A_row, A_col).read(),
+                    B.handle(B_row, B_col).rw());
+  t.set_symbol("DTRSM");
+  t->*[side, uplo, transa, diag, alpha](cudaStream_t s, auto sA, auto sB) {
+    auto& h = get_cublas_handle();
+    cuda_safe_call(cublasSetStream(h, s));
+
+    cuda_safe_call(cublasDtrsm(
+      h,
+      side,
+      uplo,
+      transa,
+      diag,
+      sB.extent(0),
+      sB.extent(1),
+      &alpha,
+      sA.data_handle(),
+      sA.stride(1),
+      sB.data_handle(),
+      sB.stride(1)));
+  };
+}
+
+void PDNRM2_HOST(matrix<double>* A, double* result)
+{
+#ifdef HAVE_DOT
+  reserved::dot::set_current_color("red");
+#endif
+
+  for (size_t rowb = 0; rowb < A->mt; rowb++)
+  {
+    for (size_t colb = 0; colb < A->nt; colb++)
+    {
+      ctx.host_launch(A->handle(rowb, colb).read())->*[=](auto sA) {
+        double res2 = 0.0;
+        for (size_t col = 0; col < sA.extent(1); col++)
+        {
+          for (size_t row = 0; row < sA.extent(0); row++)
+          {
+            double v = sA(row, col);
+            res2 += v * v;
+          }
+        }
+        *result += res2;
+      };
+    }
+  }
+}
+
+void PDPOTRF(matrix<double>& A)
+{
+  nvtx_range r("PDPOTRF");
+
+#ifdef HAVE_DOT
+  reserved::dot::set_current_color("yellow");
+#endif
+
+  assert(A.m == A.n);
+  assert(A.mt == A.nt);
+
+  int NBLOCKS = A.mt;
+  assert(A.mb == A.nb);
+
+  cuda_safe_call(cudaSetDevice(0));
+
+  for (int K = 0; K < NBLOCKS; K++)
+  {
+    int dev_akk = A.get_preferred_devid(K, K);
+    cuda_safe_call(cudaSetDevice(A.get_preferred_devid(K, K)));
+    DPOTRF(CUBLAS_FILL_MODE_LOWER, A, K, K);
+
+    for (int row = K + 1; row < NBLOCKS; row++)
+    {
+      cuda_safe_call(cudaSetDevice(A.get_preferred_devid(row, K)));
+      DTRSM(CUBLAS_SIDE_RIGHT, CUBLAS_FILL_MODE_LOWER, CUBLAS_OP_T, CUBLAS_DIAG_NON_UNIT, 1.0, A, K, K, A, row, K);
+
+      for (int col = K + 1; col < row; col++)
+      {
+        cuda_safe_call(cudaSetDevice(A.get_preferred_devid(row, col)));
+        DGEMM(CUBLAS_OP_N, CUBLAS_OP_T, -1.0, A, row, K, A, col, K, 1.0, A, row, col);
+      }
+
+      cuda_safe_call(cudaSetDevice(A.get_preferred_devid(row, row)));
+      DSYRK(CUBLAS_FILL_MODE_LOWER, CUBLAS_OP_N, -1.0, A, row, K, 1.0, A, row, row);
+    }
+  }
+  cuda_safe_call(cudaSetDevice(0));
+}
+
+// Algorithm from PLASMA
+void PDTRSM(cublasSideMode_t side,
+            cublasFillMode_t uplo,
+            cublasOperation_t trans,
+            cublasDiagType_t diag,
+            double alpha,
+            class matrix<double>& A,
+            class matrix<double>& B)
+{
+  nvtx_range r("PDTRSM");
+
+  //    std::cout << "[PDTRSM] START B MT " << B.mt << " NT " << B.nt << '\n';
+
+  if (side == CUBLAS_SIDE_LEFT)
+  {
+    if (uplo == CUBLAS_FILL_MODE_UPPER)
+    {
+      // TODO
+      assert(0);
+      abort();
+    }
+    else
+    {
+      //===========================================
+      // CUBLAS_SIDE_LEFT / CUBLAS_FILL_MODE_LOWER / CUBLAS_OP_N
+      //===========================================
+      if (trans == CUBLAS_OP_N)
+      {
+        for (size_t k = 0; k < B.mt; k++)
+        {
+          double lalpha = k == 0 ? alpha : 1.0;
+          for (size_t n = 0; n < B.nt; n++)
+          {
+            cuda_safe_call(cudaSetDevice(A.get_preferred_devid(k, k)));
+            DTRSM(side, uplo, trans, diag, lalpha, A, k, k, B, k, n);
+          }
+          for (size_t m = k + 1; m < B.mt; m++)
+          {
+            for (size_t n = 0; n < B.nt; n++)
+            {
+              cuda_safe_call(cudaSetDevice(A.get_preferred_devid(m, k)));
+              DGEMM(CUBLAS_OP_N, CUBLAS_OP_N, -1.0, A, m, k, B, k, n, lalpha, B, m, n);
+            }
+          }
+        }
+      }
+      //================================================
+      // CUBLAS_SIDE_LEFT / CUBLAS_FILL_MODE_LOWER / CUBLAS_OP_[C|T]
+      //================================================
+      else
+      {
+        for (size_t k = 0; k < B.mt; k++)
+        {
+          double lalpha = k == 0 ? alpha : 1.0;
+          for (size_t n = 0; n < B.nt; n++)
+          {
+            cuda_safe_call(cudaSetDevice(A.get_preferred_devid(B.mt - k - 1, B.mt - k - 1)));
+            DTRSM(side, uplo, trans, diag, lalpha, A, B.mt - k - 1, B.mt - k - 1, B, B.mt - k - 1, n);
+          }
+          for (size_t m = k + 1; m < B.mt; m++)
+          {
+            for (size_t n = 0; n < B.nt; n++)
+            {
+              cuda_safe_call(cudaSetDevice(A.get_preferred_devid(B.mt - k - 1, B.mt - 1 - m)));
+              DGEMM(
+                trans, CUBLAS_OP_N, -1.0, A, B.mt - k - 1, B.mt - 1 - m, B, B.mt - k - 1, n, lalpha, B, B.mt - 1 - m, n);
+            }
+          }
+        }
+      }
+    }
+  }
+  else
+  {
+    // TODO
+    abort();
+  }
+  cuda_safe_call(cudaSetDevice(0));
+  //    std::cout << "[PDTRSM] END" << '\n';
+}
+
+void PDPOTRS(matrix<double>& A, class matrix<double>& B, cublasFillMode_t uplo)
+{
+  nvtx_range r("PDPOTRS");
+
+#ifdef HAVE_DOT
+  reserved::dot::set_current_color("green");
+#endif
+
+  //    std::cout << "[PDPOTRS] START" << '\n';
+  // Call the parallel functions.
+  PDTRSM(
+    CUBLAS_SIDE_LEFT, uplo, uplo == CUBLAS_FILL_MODE_UPPER ? CUBLAS_OP_T : CUBLAS_OP_N, CUBLAS_DIAG_NON_UNIT, 1.0, A, B);
+
+#ifdef HAVE_DOT
+  reserved::dot::set_current_color("darkgreen");
+#endif
+
+  PDTRSM(
+    CUBLAS_SIDE_LEFT, uplo, uplo == CUBLAS_FILL_MODE_UPPER ? CUBLAS_OP_N : CUBLAS_OP_T, CUBLAS_DIAG_NON_UNIT, 1.0, A, B);
+  //    std::cout << "[PDPOTRS] END" << '\n';
+}
+
+/*****************************************************************************
+ * Parallel tile matrix-matrix
+ *multiplication.
+ * @see plasma_omp_dgemm
+ ******************************************************************************/
+void PDGEMM(cublasOperation_t transa,
+            cublasOperation_t transb,
+            double alpha,
+            class matrix<double>& A,
+            class matrix<double>& B,
+            double beta,
+            class matrix<double>& C)
+{
+  nvtx_range r("PDGEMM");
+
+#ifdef HAVE_DOT
+  reserved::dot::set_current_color("blue");
+#endif
+
+  for (size_t m = 0; m < C.mt; m++)
+  {
+    for (size_t n = 0; n < C.nt; n++)
+    {
+      //=========================================
+      // alpha*A*B does not contribute; scale C
+      //=========================================
+      int inner_k = transa == CUBLAS_OP_N ? A.n : A.m;
+      if (alpha == 0.0 || inner_k == 0)
+      {
+        DGEMM(transa, transb, alpha, A, 0, 0, B, 0, 0, beta, C, m, n);
+      }
+      else if (transa == CUBLAS_OP_N)
+      {
+        //================================
+        // CUBLAS_OP_N / CUBLAS_OP_N
+        //================================
+        if (transb == CUBLAS_OP_N)
+        {
+          for (size_t k = 0; k < A.nt; k++)
+          {
+            double zbeta = k == 0 ? beta : 1.0;
+            DGEMM(transa, transb, alpha, A, m, k, B, k, n, zbeta, C, m, n);
+          }
+        }
+        //=====================================
+        // CUBLAS_OP_N / CUBLAS_OP_T
+        //=====================================
+        else
+        {
+          for (size_t k = 0; k < A.nt; k++)
+          {
+            double zbeta = k == 0 ? beta : 1.0;
+            DGEMM(transa, transb, alpha, A, m, k, B, n, k, zbeta, C, m, n);
+          }
+        }
+      }
+      else
+      {
+        //=====================================
+        // CUBLAS_OP_T / CUBLAS_OP_N
+        //=====================================
+        if (transb == CUBLAS_OP_N)
+        {
+          for (size_t k = 0; k < A.mt; k++)
+          {
+            double zbeta = k == 0 ? beta : 1.0;
+            DGEMM(transa, transb, alpha, A, k, m, B, k, n, zbeta, C, m, n);
+          }
+        }
+        //==========================================
+        // CUBLAS_OP_T / CUBLAS_OP_T
+        //==========================================
+        else
+        {
+          for (size_t k = 0; k < A.mt; k++)
+          {
+            double zbeta = k == 0 ? beta : 1.0;
+            DGEMM(transa, transb, alpha, A, k, m, B, n, k, zbeta, C, m, n);
+          }
+        }
+      }
+    }
+  }
+}
+
+int main(int argc, char** argv)
+{
+  int N  = 1024;
+  int NB = 128;
+
+  if (argc > 1)
+  {
+    N = atoi(argv[1]);
+  }
+
+  if (argc > 2)
+  {
+    NB = atoi(argv[2]);
+  }
+
+  int check_result = 1;
+  if (getenv("CHECK_RESULT"))
+  {
+    check_result = atoi(getenv("CHECK_RESULT"));
+  }
+
+  assert(N % NB == 0);
+
+  // Use pools of preallocated blocks
+  auto fixed_alloc = block_allocator<fixed_size_allocator>(ctx, NB * NB * sizeof(double));
+  ctx.set_allocator(fixed_alloc);
+
+  // Set up CUBLAS and CUSOLVER
+  int ndevs;
+  cuda_safe_call(cudaGetDeviceCount(&ndevs));
+
+  for (int d = 0; d < ndevs; d++)
+  {
+    auto lX = ctx.logical_data(shape_of<slice<double>>(1));
+    ctx.parallel_for(exec_place::device(d), lX.shape(), lX.write())->*[] __device__(size_t, auto) {};
+    cuda_safe_call(cudaSetDevice(d));
+    get_cublas_handle();
+    get_cusolver_handle();
+  }
+
+  cuda_safe_call(cudaSetDevice(0));
+
+  matrix<double> A(N, N, NB, NB, true, "A");
+  matrix<double> Aref(N, N, NB, NB, false, "Aref");
+
+  // (Hilbert matrix + 2*N*Id) to have a diagonal dominant matrix
+  auto hilbert = [=] __host__ __device__(size_t row, size_t col) {
+    return 1.0 / (col + row + 1.0) + 2.0 * N * (col == row);
+  };
+
+  if (check_result)
+  {
+    Aref.fill(hilbert);
+  }
+
+  A.fill(hilbert);
+
+  /* Right-hand side */
+  matrix<double> B_potrs(N, 1, NB, 1, false, "B");
+  matrix<double> Bref_potrs(N, 1, NB, 1, false, "Bref");
+
+  if (check_result)
+  {
+    auto rhs_vals = [] __host__ __device__(size_t row, size_t /*unused*/) { return 1.0 * (row + 1); };
+    B_potrs.fill(rhs_vals);
+    Bref_potrs.fill(rhs_vals);
+  }
+
+  //    // Compute ||Bref||
+  double Bref_nrm2 = 0.0;
+  double res_nrm2  = 0.0;
+
+  if (check_result)
+  {
+    PDNRM2_HOST(&Bref_potrs, &Bref_nrm2);
+  }
+
+  cudaEvent_t startEvent_pdpotrf, stopEvent_pdpotrf;
+  float milliseconds_pdpotrf = 0;
+
+  //    for (size_t row = 0; row < A.mt; row++)
+  //    {
+  //        for (size_t col = 0; col <= row; col++)
+  //        {
+  //            cuda_safe_call(cudaSetDevice(A.get_preferred_devid(row, col)));
+  //            NOOP(A, row, col);
+  //        }
+  //    }
+
+  cuda_safe_call(cudaSetDevice(0));
+
+  cuda_safe_call(cudaStreamSynchronize(ctx.fence()));
+
+  cuda_safe_call(cudaEventCreate(&startEvent_pdpotrf));
+  cuda_safe_call(cudaEventCreate(&stopEvent_pdpotrf));
+
+  cuda_safe_call(cudaEventRecord(startEvent_pdpotrf, ctx.fence()));
+
+  PDPOTRF(A);
+
+  cuda_safe_call(cudaSetDevice(0));
+  cuda_safe_call(cudaEventRecord(stopEvent_pdpotrf, ctx.fence()));
+
+  /*
+   *  POTRS
+   */
+
+  if (check_result)
+  {
+    // Solve AX = B and put the result in B
+    PDPOTRS(A, B_potrs, CUBLAS_FILL_MODE_LOWER);
+
+    // Compute (AX - B)
+    // Bref = (Aref*B - Bref)
+    PDGEMM(CUBLAS_OP_N, CUBLAS_OP_N, 1.0, Aref, B_potrs, -1.0, Bref_potrs);
+
+    // Compute ||AX - B|| = ||Bref||
+    PDNRM2_HOST(&Bref_potrs, &res_nrm2);
+  }
+
+  ctx.finalize();
+
+  cuda_safe_call(cudaEventElapsedTime(&milliseconds_pdpotrf, startEvent_pdpotrf, stopEvent_pdpotrf));
+
+  double gflops_pdpotrf = 1.0 / 3.0 * ((double) N * (double) N * (double) N) / (1000000000.0);
+  std::cout << "[PDPOTRF] ELAPSED: " << milliseconds_pdpotrf
+            << " ms, GFLOPS: " << gflops_pdpotrf / (milliseconds_pdpotrf / 1000.0) << '\n';
+
+  if (check_result)
+  {
+    if (double residual = sqrt(res_nrm2) / sqrt(Bref_nrm2); residual >= 0.01)
+    {
+      std::cerr << "[POTRS] ||AX - B|| : " << sqrt(res_nrm2) << '\n';
+      std::cerr << "[POTRS] ||B|| : " << sqrt(Bref_nrm2) << '\n';
+      std::cerr << "[POTRS] RESIDUAL (||AX - B||/||B||) : " << residual << '\n';
+      assert(!"Algorithm did not converge.");
+    }
+  }
+}

cccl_upstream/cudax/test/stf/examples/07-cholesky-unified.cu

@@ -0,0 +1,693 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+#include <cuda/experimental/__stf/utility/nvtx.cuh>
+
+#define TILED
+
+using namespace cuda::experimental::stf;
+
+// The backend used in this example only depends on that type
+using backend_type = stream_ctx;
+// using backend_type = graph_ctx;
+
+// Global for the sake of simplicity !
+backend_type ctx;
+
+/* Get a CUBLAS handle valid on the current execution place, or initialize it lazily */
+cublasHandle_t& get_cublas_handle(const exec_place& ep = exec_place::current_device())
+{
+  static std::unordered_map<exec_place, cublasHandle_t, hash<exec_place>> cublas_handles;
+  auto& result = cublas_handles[ep];
+  if (result == cublasHandle_t())
+  { // not found, default value inserted
+    // Lazy initialization, and save the handle for future use
+    cuda_safe_call(cublasCreate(&result));
+  }
+  return result;
+}
+
+/* Get a CUSOLVER handle valid on the current execution place, or initialize it lazily */
+cusolverDnHandle_t& get_cusolver_handle(const exec_place& ep = exec_place::current_device())
+{
+  static std::unordered_map<exec_place, cusolverDnHandle_t, hash<exec_place>> cusolver_handles;
+  auto& result = cusolver_handles[ep];
+  if (result == cusolverDnHandle_t())
+  { // not found, default value inserted
+    // Lazy initialization, and save the handle for future use
+    cuda_safe_call(cusolverDnCreate(&result));
+  }
+  return result;
+}
+
+template <typename T>
+class matrix
+{
+public:
+  matrix(int NROWS, int NCOLS, int BLOCKSIZE_ROWS, int BLOCKSIZE_COLS, bool is_sym, const char* _symbol = "matrix")
+  {
+    symbol = _symbol;
+
+    sym_matrix = is_sym;
+
+    m  = NROWS;
+    mb = BLOCKSIZE_ROWS;
+
+    n  = NCOLS;
+    nb = BLOCKSIZE_COLS;
+
+    assert(m % mb == 0);
+    assert(n % nb == 0);
+
+    // cuda_safe_call(cudaMallocHost(&h_array, m*n*sizeof(T)));
+    // fprintf(stderr, "Allocating %ld x %ld x %ld = %ld bytes (%f GB) on host for %s\n", m, n, sizeof(T), s,
+    //        s / (1024.0 * 1024.0 * 1024.0), _symbol);
+    h_array.resize(m * n);
+    cuda_safe_call(cudaHostRegister(&h_array[0], h_array.size() * sizeof(T), cudaHostRegisterPortable));
+
+    // Compute the number of blocks
+    mt = m / mb;
+    nt = n / nb;
+
+    handles.resize(mt * nt);
+
+    for (size_t colb = 0; colb < nt; colb++)
+    {
+      int low_rowb = sym_matrix ? colb : 0;
+      for (size_t rowb = low_rowb; rowb < mt; rowb++)
+      {
+        T* addr_h = get_block_h(rowb, colb);
+        auto& h   = handle(rowb, colb);
+
+#ifdef TILED
+        // tiles are stored contiguously
+        size_t ld = mb;
+#else
+        size_t ld = m;
+#endif
+        std::ignore = ld; // work around bug in compiler
+        h           = ctx.logical_data(make_slice(addr_h, std::tuple{mb, nb}, ld));
+        h.set_symbol(std::string(symbol) + "_" + std::to_string(rowb) + "_" + std::to_string(colb));
+      }
+    }
+
+    cuda_safe_call(cudaGetDeviceCount(&ndevs));
+    for (int a = 1; a * a <= ndevs; a++)
+    {
+      if (ndevs % a == 0)
+      {
+        grid_p = a;
+        grid_q = ndevs / a;
+      }
+    }
+
+    assert(grid_p * grid_q == ndevs);
+
+    // std::cout << "FOUND " << ndevs << " DEVICES "
+    //          << "p=" << grid_p << " q=" << grid_q << '\n';
+  }
+
+  int get_preferred_devid(int row, int col)
+  {
+    return (row % grid_p) + (col % grid_q) * grid_p;
+  }
+
+  auto& handle(int row, int col)
+  {
+    return handles[row + col * mt];
+  }
+
+  size_t get_index(size_t row, size_t col)
+  {
+#ifdef TILED
+    // Find which tile contains this element
+    int tile_row = row / mb;
+    int tile_col = col / nb;
+
+    size_t tile_size = mb * nb;
+
+    // Look for the index of the beginning of the tile
+    size_t tile_start = (tile_row + mt * tile_col) * tile_size;
+
+    // Offset within the tile
+    size_t offset = (row % mb) + (col % nb) * mb;
+
+    return tile_start + offset;
+#else
+    return row + col * m;
+#endif
+  }
+
+  T* get_block_h(int brow, int bcol)
+  {
+    size_t index = get_index(brow * mb, bcol * nb);
+    return &h_array[index];
+  }
+
+  // Fill with func(Matrix*,row, col)
+  template <typename Fun>
+  void fill(Fun&& fun)
+  {
+    // Fill blocks by blocks
+    for (size_t colb = 0; colb < nt; colb++)
+    {
+      size_t low_rowb = sym_matrix ? colb : 0;
+      for (size_t rowb = low_rowb; rowb < mt; rowb++)
+      {
+        // Each task fills a block
+        ctx.host_launch(handle(rowb, colb).write())->*[=, self = this](auto sA) {
+          for (size_t lcol = 0; lcol < sA.extent(1); lcol++)
+          {
+            size_t col = lcol + colb * sA.extent(1);
+            for (size_t lrow = 0; lrow < sA.extent(0); lrow++)
+            {
+              size_t row     = lrow + rowb * sA.extent(0);
+              sA(lrow, lcol) = fun(*self, row, col);
+            }
+          }
+        };
+      }
+    }
+  }
+
+  std::vector<T> h_array;
+  size_t m; // nrows
+  size_t n; // ncols
+
+  // Is this a sym matrix ? (lower assumed)
+  bool sym_matrix;
+
+  size_t mb; // block size (rows)
+  size_t nb; // block size (cols)
+
+  size_t mt; // number of column blocks
+  size_t nt; // number of row blocks
+
+  // abstract data handles
+  std::vector<logical_data<slice<double, 2>>> handles;
+
+  const char* symbol;
+
+  // for the mapping
+  int ndevs;
+  int grid_p, grid_q;
+};
+
+void DPOTRF(cublasFillMode_t uplo, class matrix<double>& A, int A_row, int A_col)
+{
+  auto& Akk    = A.handle(A_row, A_col);
+  size_t m_akk = Akk.shape().extent(0);
+  // Note that the handle may be different from the actual handle...
+  int Lwork_expected;
+  cuda_safe_call(cusolverDnDpotrf_bufferSize(get_cusolver_handle(), uplo, m_akk, nullptr, 0, &Lwork_expected));
+
+  auto potrf_buffer = ctx.logical_data<double>(Lwork_expected);
+  auto devInfo      = ctx.logical_data(shape_of<slice<int>>(1));
+
+  auto t = ctx.task(Akk.rw(), potrf_buffer.write(), devInfo.write());
+  //    t.set_symbol("DPOTRF");
+  t->*[&](cudaStream_t s, auto sAkk, auto buffer, auto info) {
+    auto& h = get_cusolver_handle();
+    cuda_safe_call(cusolverDnSetStream(h, s));
+
+    cuda_safe_call(cusolverDnDpotrf(
+      h,
+      uplo,
+      sAkk.extent(0),
+      sAkk.data_handle(),
+      sAkk.stride(1),
+      buffer.data_handle(),
+      buffer.extent(0),
+      info.data_handle()));
+  };
+}
+
+void DGEMM(
+  cublasOperation_t transa,
+  cublasOperation_t transb,
+  double alpha,
+  class matrix<double>& A,
+  int A_row,
+  int A_col,
+  class matrix<double>& B,
+  int B_row,
+  int B_col,
+  double beta,
+  class matrix<double>& C,
+  int C_row,
+  int C_col)
+{
+  auto ignored = get_cublas_handle();
+  auto t       = ctx.task(A.handle(A_row, A_col).read(), B.handle(B_row, B_col).read(), C.handle(C_row, C_col).rw());
+  //    t.set_symbol("DGEMM");
+  t->*[&](cudaStream_t s, auto sA, auto sB, auto sC) {
+    auto& h = get_cublas_handle();
+    cuda_safe_call(cublasSetStream(h, s));
+
+    auto k = (transa == CUBLAS_OP_N) ? sA.extent(1) : sA.extent(0);
+    cuda_safe_call(cublasDgemm(
+      h,
+      transa,
+      transb,
+      sC.extent(0),
+      sC.extent(1),
+      k,
+      &alpha,
+      sA.data_handle(),
+      sA.stride(1),
+      sB.data_handle(),
+      sB.stride(1),
+      &beta,
+      sC.data_handle(),
+      sC.stride(1)));
+  };
+}
+
+void DSYRK(
+  cublasFillMode_t uplo,
+  cublasOperation_t trans,
+  double alpha,
+  class matrix<double>& A,
+  int A_row,
+  int A_col,
+  double beta,
+  class matrix<double>& C,
+  int C_row,
+  int C_col)
+{
+  auto ignored = get_cublas_handle();
+  auto t       = ctx.task(A.handle(A_row, A_col).read(), C.handle(C_row, C_col).rw());
+  //    t.set_symbol("DSYRK");
+  t->*[&](cudaStream_t s, auto sA, auto sC) {
+    auto& h = get_cublas_handle();
+    cuda_safe_call(cublasSetStream(h, s));
+
+    // number of rows of matrix op(A) and C
+    auto n = sC.extent(0);
+
+    // number of columns of matrix op(A)
+    auto k = (trans == CUBLAS_OP_N) ? sA.extent(1) : sA.extent(0);
+
+    cuda_safe_call(
+      cublasDsyrk(h, uplo, trans, n, k, &alpha, sA.data_handle(), sA.stride(1), &beta, sC.data_handle(), sC.stride(1)));
+  };
+}
+
+void DTRSM(
+  cublasSideMode_t side,
+  cublasFillMode_t uplo,
+  cublasOperation_t transa,
+  cublasDiagType_t diag,
+  double alpha,
+  class matrix<double>& A,
+  int A_row,
+  int A_col,
+  class matrix<double>& B,
+  int B_row,
+  int B_col)
+{
+  auto ignored = get_cublas_handle();
+  auto t       = ctx.task(A.handle(A_row, A_col).read(), B.handle(B_row, B_col).rw());
+  //    t.set_symbol("DTRSM");
+  t->*[&](cudaStream_t s, auto sA, auto sB) {
+    auto& h = get_cublas_handle();
+    cuda_safe_call(cublasSetStream(h, s));
+
+    cuda_safe_call(cublasDtrsm(
+      h,
+      side,
+      uplo,
+      transa,
+      diag,
+      sB.extent(0),
+      sB.extent(1),
+      &alpha,
+      sA.data_handle(),
+      sA.stride(1),
+      sB.data_handle(),
+      sB.stride(1)));
+  };
+}
+
+void PDNRM2_HOST(matrix<double>* A, double* result)
+{
+#ifdef HAVE_DOT
+  reserved::dot::set_current_color("red");
+#endif
+
+  for (size_t rowb = 0; rowb < A->mt; rowb++)
+  {
+    for (size_t colb = 0; colb < A->nt; colb++)
+    {
+      ctx.host_launch(A->handle(rowb, colb).read())->*[=](auto sA) {
+        double res2 = 0.0;
+        for (size_t col = 0; col < sA.extent(1); col++)
+        {
+          for (size_t row = 0; row < sA.extent(0); row++)
+          {
+            double v = sA(row, col);
+            res2 += v * v;
+          }
+        }
+        *result += res2;
+      };
+    }
+  }
+}
+
+void PDPOTRF(matrix<double>& A)
+{
+  nvtx_range r("PDPOTRF");
+
+#ifdef HAVE_DOT
+  reserved::dot::set_current_color("yellow");
+#endif
+
+  assert(A.m == A.n);
+  assert(A.mt == A.nt);
+
+  int NBLOCKS = A.mt;
+  assert(A.mb == A.nb);
+
+  cuda_safe_call(cudaSetDevice(0));
+
+  for (int K = 0; K < NBLOCKS; K++)
+  {
+    int dev_akk = A.get_preferred_devid(K, K);
+    cuda_safe_call(cudaSetDevice(A.get_preferred_devid(K, K)));
+    DPOTRF(CUBLAS_FILL_MODE_LOWER, A, K, K);
+
+    for (int row = K + 1; row < NBLOCKS; row++)
+    {
+      cuda_safe_call(cudaSetDevice(A.get_preferred_devid(row, K)));
+      DTRSM(CUBLAS_SIDE_RIGHT, CUBLAS_FILL_MODE_LOWER, CUBLAS_OP_T, CUBLAS_DIAG_NON_UNIT, 1.0, A, K, K, A, row, K);
+
+      for (int col = K + 1; col < row; col++)
+      {
+        cuda_safe_call(cudaSetDevice(A.get_preferred_devid(row, col)));
+        DGEMM(CUBLAS_OP_N, CUBLAS_OP_T, -1.0, A, row, K, A, col, K, 1.0, A, row, col);
+      }
+
+      cuda_safe_call(cudaSetDevice(A.get_preferred_devid(row, row)));
+      DSYRK(CUBLAS_FILL_MODE_LOWER, CUBLAS_OP_N, -1.0, A, row, K, 1.0, A, row, row);
+    }
+  }
+  cuda_safe_call(cudaSetDevice(0));
+}
+
+// Algorithm from PLASMA
+void PDTRSM(cublasSideMode_t side,
+            cublasFillMode_t uplo,
+            cublasOperation_t trans,
+            cublasDiagType_t diag,
+            double alpha,
+            class matrix<double>& A,
+            class matrix<double>& B)
+{
+  nvtx_range r("PDTRSM");
+
+  //    std::cout << "[PDTRSM] START B MT " << B.mt << " NT " << B.nt << '\n';
+
+  if (side == CUBLAS_SIDE_LEFT)
+  {
+    if (uplo == CUBLAS_FILL_MODE_UPPER)
+    {
+      // TODO
+      assert(0);
+      abort();
+    }
+    else
+    {
+      //===========================================
+      // CUBLAS_SIDE_LEFT / CUBLAS_FILL_MODE_LOWER / CUBLAS_OP_N
+      //===========================================
+      if (trans == CUBLAS_OP_N)
+      {
+        for (size_t k = 0; k < B.mt; k++)
+        {
+          double lalpha = k == 0 ? alpha : 1.0;
+          for (size_t n = 0; n < B.nt; n++)
+          {
+            cuda_safe_call(cudaSetDevice(A.get_preferred_devid(k, k)));
+            DTRSM(side, uplo, trans, diag, lalpha, A, k, k, B, k, n);
+          }
+          for (size_t m = k + 1; m < B.mt; m++)
+          {
+            for (size_t n = 0; n < B.nt; n++)
+            {
+              cuda_safe_call(cudaSetDevice(A.get_preferred_devid(m, k)));
+              DGEMM(CUBLAS_OP_N, CUBLAS_OP_N, -1.0, A, m, k, B, k, n, lalpha, B, m, n);
+            }
+          }
+        }
+      }
+      //================================================
+      // CUBLAS_SIDE_LEFT / CUBLAS_FILL_MODE_LOWER / CUBLAS_OP_[C|T]
+      //================================================
+      else
+      {
+        for (size_t k = 0; k < B.mt; k++)
+        {
+          double lalpha = k == 0 ? alpha : 1.0;
+          for (size_t n = 0; n < B.nt; n++)
+          {
+            cuda_safe_call(cudaSetDevice(A.get_preferred_devid(B.mt - k - 1, B.mt - k - 1)));
+            DTRSM(side, uplo, trans, diag, lalpha, A, B.mt - k - 1, B.mt - k - 1, B, B.mt - k - 1, n);
+          }
+          for (size_t m = k + 1; m < B.mt; m++)
+          {
+            for (size_t n = 0; n < B.nt; n++)
+            {
+              cuda_safe_call(cudaSetDevice(A.get_preferred_devid(B.mt - k - 1, B.mt - 1 - m)));
+              DGEMM(
+                trans, CUBLAS_OP_N, -1.0, A, B.mt - k - 1, B.mt - 1 - m, B, B.mt - k - 1, n, lalpha, B, B.mt - 1 - m, n);
+            }
+          }
+        }
+      }
+    }
+  }
+  else
+  {
+    // TODO
+    abort();
+  }
+  cuda_safe_call(cudaSetDevice(0));
+  //    std::cout << "[PDTRSM] END" << '\n';
+}
+
+void PDPOTRS(matrix<double>& A, class matrix<double>& B, cublasFillMode_t uplo)
+{
+  nvtx_range r("PDPOTRS");
+
+#ifdef HAVE_DOT
+  reserved::dot::set_current_color("green");
+#endif
+
+  //    std::cout << "[PDPOTRS] START" << '\n';
+  // Call the parallel functions.
+  PDTRSM(
+    CUBLAS_SIDE_LEFT, uplo, uplo == CUBLAS_FILL_MODE_UPPER ? CUBLAS_OP_T : CUBLAS_OP_N, CUBLAS_DIAG_NON_UNIT, 1.0, A, B);
+
+#ifdef HAVE_DOT
+  reserved::dot::set_current_color("darkgreen");
+#endif
+
+  PDTRSM(
+    CUBLAS_SIDE_LEFT, uplo, uplo == CUBLAS_FILL_MODE_UPPER ? CUBLAS_OP_N : CUBLAS_OP_T, CUBLAS_DIAG_NON_UNIT, 1.0, A, B);
+  //    std::cout << "[PDPOTRS] END" << '\n';
+}
+
+/*****************************************************************************
+ * Parallel tile matrix-matrix
+ *multiplication.
+ * @see plasma_omp_dgemm
+ ******************************************************************************/
+void PDGEMM(cublasOperation_t transa,
+            cublasOperation_t transb,
+            double alpha,
+            class matrix<double>& A,
+            class matrix<double>& B,
+            double beta,
+            class matrix<double>& C)
+{
+  nvtx_range r("PDGEMM");
+
+#ifdef HAVE_DOT
+  reserved::dot::set_current_color("blue");
+#endif
+
+  for (size_t m = 0; m < C.mt; m++)
+  {
+    for (size_t n = 0; n < C.nt; n++)
+    {
+      //=========================================
+      // alpha*A*B does not contribute; scale C
+      //=========================================
+      int inner_k = transa == CUBLAS_OP_N ? A.n : A.m;
+      if (alpha == 0.0 || inner_k == 0)
+      {
+        DGEMM(transa, transb, alpha, A, 0, 0, B, 0, 0, beta, C, m, n);
+      }
+      else if (transa == CUBLAS_OP_N)
+      {
+        //================================
+        // CUBLAS_OP_N / CUBLAS_OP_N
+        //================================
+        if (transb == CUBLAS_OP_N)
+        {
+          for (size_t k = 0; k < A.nt; k++)
+          {
+            double zbeta = k == 0 ? beta : 1.0;
+            DGEMM(transa, transb, alpha, A, m, k, B, k, n, zbeta, C, m, n);
+          }
+        }
+        //=====================================
+        // CUBLAS_OP_N / CUBLAS_OP_T
+        //=====================================
+        else
+        {
+          for (size_t k = 0; k < A.nt; k++)
+          {
+            double zbeta = k == 0 ? beta : 1.0;
+            DGEMM(transa, transb, alpha, A, m, k, B, n, k, zbeta, C, m, n);
+          }
+        }
+      }
+      else
+      {
+        //=====================================
+        // CUBLAS_OP_T / CUBLAS_OP_N
+        //=====================================
+        if (transb == CUBLAS_OP_N)
+        {
+          for (size_t k = 0; k < A.mt; k++)
+          {
+            double zbeta = k == 0 ? beta : 1.0;
+            DGEMM(transa, transb, alpha, A, k, m, B, k, n, zbeta, C, m, n);
+          }
+        }
+        //==========================================
+        // CUBLAS_OP_T / CUBLAS_OP_T
+        //==========================================
+        else
+        {
+          for (size_t k = 0; k < A.mt; k++)
+          {
+            double zbeta = k == 0 ? beta : 1.0;
+            DGEMM(transa, transb, alpha, A, k, m, B, n, k, zbeta, C, m, n);
+          }
+        }
+      }
+    }
+  }
+}
+
+int main(int argc, char** argv)
+{
+  int N  = 1024;
+  int NB = 128;
+
+  if (argc > 1)
+  {
+    N = atoi(argv[1]);
+  }
+
+  if (argc > 2)
+  {
+    NB = atoi(argv[2]);
+  }
+
+  assert(N % NB == 0);
+
+  // Set up CUBLAS and CUSOLVER
+  int ndevs;
+  cuda_safe_call(cudaGetDeviceCount(&ndevs));
+
+  cuda_safe_call(cudaSetDevice(0));
+
+  matrix<double> A(N, N, NB, NB, true, "A");
+  matrix<double> Aref(N, N, NB, NB, false, "Aref");
+
+  // (Hilbert matrix + 2*N*Id) to have a diagonal dominant matrix
+  auto hilbert = [](matrix<double>& mat, int row, int col) {
+    return 1.0 / (col + row + 1.0) + 2.0 * mat.n * (col == row);
+  };
+
+  Aref.fill(hilbert);
+  A.fill(hilbert);
+
+  /* Right-hand side */
+  matrix<double> B_potrs(N, 1, NB, 1, false, "B");
+  matrix<double> Bref_potrs(N, 1, NB, 1, false, "Bref");
+
+  auto rhs_vals = [](matrix<double>& /*unused*/, int row, int /*unused*/) {
+    return 1.0 * (row + 1);
+  };
+  B_potrs.fill(rhs_vals);
+  Bref_potrs.fill(rhs_vals);
+
+  int check_result = 1;
+  if (getenv("CHECK_RESULT"))
+  {
+    check_result = atoi(getenv("CHECK_RESULT"));
+  }
+
+  //    // Compute ||Bref||
+  double Bref_nrm2 = 0.0;
+  double res_nrm2  = 0.0;
+
+  if (check_result)
+  {
+    PDNRM2_HOST(&Bref_potrs, &Bref_nrm2);
+  }
+
+  //    for (size_t row = 0; row < A.mt; row++)
+  //    {
+  //        for (size_t col = 0; col <= row; col++)
+  //        {
+  //            cuda_safe_call(cudaSetDevice(A.get_preferred_devid(row, col)));
+  //            NOOP(A, row, col);
+  //        }
+  //    }
+
+  PDPOTRF(A);
+
+  /*
+   *  POTRS
+   */
+
+  if (check_result)
+  {
+    // Solve AX = B and put the result in B
+    PDPOTRS(A, B_potrs, CUBLAS_FILL_MODE_LOWER);
+
+    // Compute (AX - B)
+    // Bref = (Aref*B - Bref)
+    PDGEMM(CUBLAS_OP_N, CUBLAS_OP_N, 1.0, Aref, B_potrs, -1.0, Bref_potrs);
+
+    // Compute ||AX - B|| = ||Bref||
+    PDNRM2_HOST(&Bref_potrs, &res_nrm2);
+  }
+
+  ctx.finalize();
+
+  if (check_result)
+  {
+    double residual = sqrt(res_nrm2) / sqrt(Bref_nrm2);
+    // std::cout << "[POTRS] ||AX - B|| : " << sqrt(res_nrm2) << '\n';
+    // std::cout << "[POTRS] ||B|| : " << sqrt(Bref_nrm2) << '\n';
+    // std::cout << "[POTRS] RESIDUAL (||AX - B||/||B||) : " << residual << '\n';
+    assert(residual < 0.01);
+  }
+
+  return 0;
+}

cccl_upstream/cudax/test/stf/examples/09-nbody-algorithm.cu

@@ -0,0 +1,129 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/experimental/stf.cuh>
+
+#include <random>
+
+using namespace cuda::experimental::stf;
+
+struct body
+{
+  // mass
+  double mass;
+  // position
+  double pos[3];
+  // speed
+  double vel[3];
+  // acceleration
+  double acc[3];
+};
+
+int main()
+{
+  constexpr double kSofteningSquared = 1e-3;
+  constexpr double kG                = 6.67259e-11;
+
+  size_t BODY_CNT = 4096;
+
+  double dt    = 0.1;
+  size_t NITER = 25;
+
+  context ctx;
+
+  std::vector<body> particles;
+  particles.resize(BODY_CNT);
+
+  // Initialize particles
+  std::random_device rd;
+  std::mt19937 gen(rd());
+  std::uniform_real_distribution<> dis(-1.0, 1.0);
+  for (auto& p : particles)
+  {
+    p.mass = 1.0;
+
+    p.pos[0] = dis(gen);
+    p.pos[1] = dis(gen);
+    p.pos[2] = dis(gen);
+
+    p.vel[0] = dis(gen);
+    p.vel[1] = dis(gen);
+    p.vel[2] = dis(gen);
+
+    p.acc[0] = 0.0;
+    p.acc[1] = 0.0;
+    p.acc[2] = 0.0;
+  }
+
+  auto h_particles = ctx.logical_data(make_slice(&particles[0], BODY_CNT));
+
+  auto fn = [dt](context ctx, logical_data<slice<body>> h_particles) {
+    // Compute accelerations
+    ctx.parallel_for(h_particles.shape(), h_particles.rw())->*[=] _CCCL_DEVICE __host__(size_t i, slice<body> p) {
+      double acc[3];
+      for (size_t k = 0; k < 3; k++)
+      {
+        acc[k] = p(i).acc[k];
+      }
+
+      for (size_t j = 0; j < p.extent(0); j++)
+      {
+        if (i != j)
+        {
+          double d[3];
+          for (size_t k = 0; k < 3; k++)
+          {
+            d[k] = p(j).pos[k] - p(i).pos[k];
+          }
+
+          double dist     = d[0] * d[0] + d[1] * d[1] + d[2] * d[2] + kSofteningSquared;
+          double dist_inv = 1.0 / sqrt(dist);
+
+          for (size_t k = 0; k < 3; k++)
+          {
+            acc[k] += d[k] * kG * p(j).mass * dist_inv * dist_inv * dist_inv;
+          }
+        }
+      }
+
+      for (size_t k = 0; k < 3; k++)
+      {
+        p(i).acc[k] = acc[k];
+      }
+    };
+
+    // Update velocity and positions
+    ctx.parallel_for(h_particles.shape(), h_particles.rw())->*[=] __host__ __device__(size_t i, slice<body> p) {
+      for (size_t k = 0; k < 3; k++)
+      {
+        p(i).vel[k] += p(i).acc[k] * dt;
+      }
+
+      for (size_t k = 0; k < 3; k++)
+      {
+        p(i).pos[k] += p(i).vel[k] * dt;
+      }
+
+      for (size_t k = 0; k < 3; k++)
+      {
+        p(i).acc[k] = 0.0;
+      }
+    };
+  };
+
+  algorithm one_iter;
+  for (size_t iter = 0; iter < NITER; iter++)
+  {
+    // fprintf(stderr, "ITER %ld\n", iter);
+    one_iter.run_as_task(fn, ctx, h_particles.rw());
+  }
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/examples/09-nbody-blocked.cu

@@ -0,0 +1,283 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/experimental/stf.cuh>
+
+#include <random>
+
+using namespace cuda::experimental::stf;
+
+struct body
+{
+  // mass
+  double mass;
+  // position
+  double pos[3];
+  // speed
+  double vel[3];
+};
+
+// Function to write VTK file for a single time step
+void writeVTKFile(context& ctx,
+                  const std::string& filename,
+                  size_t BLOCK_SIZE,
+                  size_t BODY_CNT,
+                  std::vector<logical_data<slice<body>>> parts)
+{
+  std::ofstream outfile(filename);
+
+  if (!outfile)
+  {
+    std::cerr << "Error opening file: " << filename << '\n';
+    return;
+  }
+
+  outfile << "# vtk DataFile Version 4.2\n";
+  outfile << "Position Data\n";
+  outfile << "ASCII\n";
+  outfile << "DATASET UNSTRUCTURED_GRID\n";
+  outfile << "POINTS " << BODY_CNT << " float\n";
+
+  std::vector<double> dump(3 * BODY_CNT);
+
+  for (size_t b = 0; b < parts.size(); b++)
+  {
+    ctx.task(exec_place::host(), parts[b].read())->*[&](cudaStream_t s, slice<const body> p) {
+      cuda_safe_call(cudaStreamSynchronize(s));
+      for (size_t i = 0; i < p.size(); i++)
+      {
+        for (size_t k = 0; k < 3; k++)
+        {
+          dump[3 * (i + b * BLOCK_SIZE) + k] = p(i).pos[k];
+        }
+      }
+    };
+  }
+
+  for (size_t p = 0; p < BODY_CNT; p++)
+  {
+    outfile << dump[3 * p] << " " << dump[3 * p + 1] << " " << dump[3 * p + 2] << "\n";
+  }
+
+  outfile.close();
+}
+
+void load_input_file(std::string filename, std::vector<body>& particles)
+{
+  std::ifstream infile(filename);
+  if (!infile)
+  {
+    std::cerr << "Error opening file." << '\n';
+    abort();
+    return;
+  }
+
+  double mass, posX, posY, posZ, velX, velY, velZ;
+  size_t ind = 0;
+
+  // Loop until we reach the end of the file
+  while (infile >> mass >> posX >> posY >> posZ >> velX >> velY >> velZ)
+  {
+    body p;
+    p.mass = mass;
+
+    p.pos[0] = posX;
+    p.pos[1] = posY;
+    p.pos[2] = posZ;
+
+    p.vel[0] = velX;
+    p.vel[1] = velY;
+    p.vel[2] = velZ;
+
+    // // Display first bodies
+    // if (ind < 10) {
+    //    fprintf(stderr, "body xyz %e %e %e dxyz %e %e %e m %e\n", p.pos[0], p.pos[1], p.pos[2], p.vel[0],
+    //            p.vel[1], p.vel[2], p.mass);
+    //}
+
+    ind++;
+    particles.push_back(p);
+  }
+
+  fprintf(stderr, "Loaded %zu bodies from %s...\n", ind, filename.c_str());
+}
+
+int main(int argc, char** argv)
+{
+  constexpr double kSofteningSquared = 1e-9;
+  //    constexpr double kG = 6.67259e-11;
+  constexpr double kG = 1.0;
+
+  size_t BODY_CNT   = 128ULL * 1024ULL;
+  size_t BLOCK_SIZE = 16 * 1024ULL;
+
+  std::vector<body> particles;
+
+  // Initialize particles
+  if (argc > 1)
+  {
+    // Get dataset from file
+    std::string filename = argv[1];
+
+    load_input_file(filename, particles);
+
+    BODY_CNT   = particles.size();
+    BLOCK_SIZE = (BODY_CNT + 7) / 8;
+  }
+  else
+  {
+    // Random distribution
+    BODY_CNT = 32ULL * 1024ULL;
+    particles.resize(BODY_CNT);
+
+    std::random_device rd;
+    std::mt19937 gen(rd());
+    std::uniform_real_distribution<> dis(-1.0, 1.0);
+    for (auto& p : particles)
+    {
+      p.mass = 1.0;
+
+      p.pos[0] = dis(gen);
+      p.pos[1] = dis(gen);
+      p.pos[2] = dis(gen);
+
+      p.vel[0] = dis(gen);
+      p.vel[1] = dis(gen);
+      p.vel[2] = dis(gen);
+    }
+  }
+
+  cuda_safe_call(cudaHostRegister(&particles[0], BODY_CNT * sizeof(body), cudaHostRegisterPortable));
+
+  double dt    = 0.005;
+  size_t NITER = 7; // 7000;
+
+  context ctx;
+
+  std::vector<logical_data<slice<body>>> parts;
+
+  // Accelerations
+  std::vector<logical_data<slice<double, 2>>> acc_parts;
+
+  size_t block_cnt = (BODY_CNT + BLOCK_SIZE - 1) / BLOCK_SIZE;
+  for (size_t i = 0; i < block_cnt; i++)
+  {
+    size_t first = i * BLOCK_SIZE;
+    size_t last  = std::min((i + 1) * BLOCK_SIZE, BODY_CNT);
+    auto p_i     = ctx.logical_data(make_slice(&particles[first], last - first));
+    parts.push_back(p_i);
+
+    auto acc_p_i = ctx.logical_data(shape_of<slice<double, 2>>(last - first, 3));
+    acc_parts.push_back(acc_p_i);
+  }
+
+  int ngpus;
+  cuda_safe_call(cudaGetDeviceCount(&ngpus));
+
+  cudaEvent_t start;
+  cuda_safe_call(cudaEventCreate(&start));
+  cuda_safe_call(cudaEventRecord(start, ctx.fence()));
+
+  for (size_t iter = 0; iter < NITER; iter++)
+  {
+    // Initialize acceleration to 0
+    for (size_t b = 0; b < block_cnt; b++)
+    {
+      ctx.launch(exec_place::device(b % ngpus), acc_parts[b].write())
+          //.set_symbol("init_acc")
+          ->*[=] _CCCL_DEVICE(auto t, slice<double, 2> acc) {
+                for (size_t i = t.rank(); i < acc.extent(0); i += t.size())
+                {
+                  for (size_t k = 0; k < 3; k++)
+                  {
+                    acc(i, k) = 0.0;
+                  }
+                }
+              };
+    }
+
+    // Compute accelerations
+    for (size_t b = 0; b < block_cnt; b++)
+    {
+      for (size_t b_other = 0; b_other < block_cnt; b_other++)
+      {
+        ctx.launch(exec_place::device(b % ngpus), parts[b].read(), parts[b_other].read(), acc_parts[b].rw())
+            //.set_symbol("compute_acc")
+            ->*[=] _CCCL_DEVICE(auto t, slice<const body> p, slice<const body> p_other, slice<double, 2> acc) {
+                  for (size_t i = t.rank(); i < p.extent(0); i += t.size())
+                  {
+                    for (size_t j = 0; j < p_other.extent(0); j++)
+                    {
+                      if ((b * BLOCK_SIZE + i) != (b_other * BLOCK_SIZE + j))
+                      {
+                        double d[3];
+                        for (size_t k = 0; k < 3; k++)
+                        {
+                          d[k] = p_other(j).pos[k] - p(i).pos[k];
+                        }
+
+                        double dist     = d[0] * d[0] + d[1] * d[1] + d[2] * d[2] + kSofteningSquared;
+                        double dist_inv = 1.0 / sqrt(dist);
+
+                        for (size_t k = 0; k < 3; k++)
+                        {
+                          acc(i, k) += d[k] * kG * p_other(j).mass * dist_inv * dist_inv * dist_inv;
+                        }
+                      }
+                    }
+                  }
+                };
+      }
+    }
+
+    for (size_t b = 0; b < block_cnt; b++)
+    {
+      // Update velocity and positions
+      ctx.launch(exec_place::device(b % ngpus), parts[b].rw(), acc_parts[b].read())
+          //.set_symbol("update")
+          ->*[=] _CCCL_DEVICE(auto t, slice<body> p, slice<const double, 2> acc) {
+                for (size_t i = t.rank(); i < p.extent(0); i += t.size())
+                {
+                  for (size_t k = 0; k < 3; k++)
+                  {
+                    p(i).vel[k] += acc(i, k) * dt;
+                  }
+
+                  for (size_t k = 0; k < 3; k++)
+                  {
+                    p(i).pos[k] += p(i).vel[k] * dt;
+                  }
+                }
+              };
+    }
+
+    // Write the VTK file for this time step
+    const char* dump_freq_str = getenv("DUMP_FREQ");
+    if (dump_freq_str && iter % atoi(dump_freq_str) == 0)
+    {
+      std::string filename = "time_step_" + std::to_string(iter) + ".vtk";
+      writeVTKFile(ctx, filename, BLOCK_SIZE, BODY_CNT, parts);
+    }
+  }
+
+  cudaEvent_t stop;
+  cuda_safe_call(cudaEventCreate(&stop));
+  cuda_safe_call(cudaEventRecord(stop, ctx.fence()));
+
+  ctx.finalize();
+
+  float elapsed;
+  cuda_safe_call(cudaEventElapsedTime(&elapsed, start, stop));
+
+  // rough approximation !
+  double FLOP_COUNT = 21.0 * (1.0 * BODY_CNT) * (1.0 * BODY_CNT) * NITER;
+
+  printf("NBODY: elapsed %f ms, %f GFLOPS\n", elapsed, FLOP_COUNT / elapsed / 1000000.0);
+}

cccl_upstream/cudax/test/stf/examples/09-nbody.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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/experimental/stf.cuh>
+
+#include <random>
+
+using namespace cuda::experimental::stf;
+
+struct body
+{
+  // mass
+  double mass;
+  // position
+  double pos[3];
+  // speed
+  double vel[3];
+  // acceleration
+  double acc[3];
+};
+
+int main()
+{
+  constexpr double kSofteningSquared = 1e-3;
+  constexpr double kG                = 6.67259e-11;
+
+  size_t BODY_CNT = 4096;
+
+  double dt    = 0.1;
+  size_t NITER = 25;
+
+  context ctx = graph_ctx();
+
+  std::vector<body> particles;
+  particles.resize(BODY_CNT);
+
+  // Initialize particles
+  std::random_device rd;
+  std::mt19937 gen(rd());
+  std::uniform_real_distribution<> dis(-1.0, 1.0);
+  for (auto& p : particles)
+  {
+    p.mass = 1.0;
+
+    p.pos[0] = dis(gen);
+    p.pos[1] = dis(gen);
+    p.pos[2] = dis(gen);
+
+    p.vel[0] = dis(gen);
+    p.vel[1] = dis(gen);
+    p.vel[2] = dis(gen);
+
+    p.acc[0] = 0.0;
+    p.acc[1] = 0.0;
+    p.acc[2] = 0.0;
+  }
+
+  auto h_particles = ctx.logical_data(make_slice(&particles[0], BODY_CNT));
+
+  ctx.repeat(NITER)->*[&](context ctx, size_t) {
+    // Compute accelerations
+    ctx.parallel_for(h_particles.shape(), h_particles.rw())->*[=] _CCCL_DEVICE __host__(size_t i, slice<body> p) {
+      double acc[3];
+      for (size_t k = 0; k < 3; k++)
+      {
+        acc[k] = p(i).acc[k];
+      }
+
+      for (size_t j = 0; j < p.extent(0); j++)
+      {
+        if (i != j)
+        {
+          double d[3];
+          for (size_t k = 0; k < 3; k++)
+          {
+            d[k] = p(j).pos[k] - p(i).pos[k];
+          }
+
+          double dist     = d[0] * d[0] + d[1] * d[1] + d[2] * d[2] + kSofteningSquared;
+          double dist_inv = 1.0 / sqrt(dist);
+
+          for (size_t k = 0; k < 3; k++)
+          {
+            acc[k] += d[k] * kG * p(j).mass * dist_inv * dist_inv * dist_inv;
+          }
+        }
+      }
+
+      for (size_t k = 0; k < 3; k++)
+      {
+        p(i).acc[k] = acc[k];
+      }
+    };
+
+    // Update velocity and positions
+    ctx.parallel_for(h_particles.shape(), h_particles.rw())->*[=] __host__ __device__(size_t i, slice<body> p) {
+      for (size_t k = 0; k < 3; k++)
+      {
+        p(i).vel[k] += p(i).acc[k] * dt;
+      }
+
+      for (size_t k = 0; k < 3; k++)
+      {
+        p(i).pos[k] += p(i).vel[k] * dt;
+      }
+
+      for (size_t k = 0; k < 3; k++)
+      {
+        p(i).acc[k] = 0.0;
+      }
+    };
+  };
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/examples/cuda_kernels_driver.cu

@@ -0,0 +1,123 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ *
+ * @brief Test that the cuda_kernel construct works with global kernels, CUfunction and CUkernel entries.
+ *
+ */
+
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+__global__ void axpy(double a, slice<const double> x, slice<double> y)
+{
+  int tid      = blockIdx.x * blockDim.x + threadIdx.x;
+  int nthreads = gridDim.x * blockDim.x;
+
+  for (int i = tid; i < x.size(); i += nthreads)
+  {
+    y(i) += a * x(i);
+  }
+}
+
+double X0(int i)
+{
+  return sin((double) i);
+}
+
+double Y0(int i)
+{
+  return cos((double) i);
+}
+
+void test(bool is_graph)
+{
+  context ctx;
+  if (is_graph)
+  {
+    ctx = graph_ctx();
+  }
+
+  const size_t N = 16;
+  double X[N], Y[N];
+
+  for (size_t i = 0; i < N; i++)
+  {
+    X[i] = X0(i);
+    Y[i] = Y0(i);
+  }
+
+  // Number of times we have applied the axpy kernel
+  int num_axpy = 0;
+
+  double alpha = 3.14;
+
+  auto lX = ctx.logical_data(X);
+  auto lY = ctx.logical_data(Y);
+
+  // runtime global kernel
+  ctx.cuda_kernel(lX.read(), lY.rw())->*[&](auto dX, auto dY) {
+    // axpy<<<16, 128, 0, ...>>>(alpha, dX, dY)
+    return cuda_kernel_desc{axpy, 16, 128, 0, alpha, dX, dY};
+  };
+  num_axpy++;
+
+  // CUfunction driver API
+  CUfunction axpy_fun;
+  cuda_safe_call(cudaGetFuncBySymbol(&axpy_fun, (void*) axpy));
+
+  ctx.cuda_kernel(lX.read(), lY.rw())->*[&](auto dX, auto dY) {
+    return cuda_kernel_desc{axpy_fun, 16, 128, 0, alpha, dX, dY};
+  };
+  num_axpy++;
+
+#if _CCCL_CTK_AT_LEAST(12, 1)
+  // CUkernel driver API
+  CUkernel axpy_kernel;
+  cuda_safe_call(cudaGetKernel(&axpy_kernel, (void*) axpy));
+
+  ctx.cuda_kernel(lX.read(), lY.rw())->*[&](auto dX, auto dY) {
+    return cuda_kernel_desc{axpy_kernel, 16, 128, 0, alpha, dX, dY};
+  };
+  num_axpy++;
+#endif
+
+  /* Some extra sanity checks, we put this in a dummy task to get access to dX and dY values */
+  ctx.task(lX.read(), lY.rw())->*[&](auto, auto dX, auto dY) {
+    int nregs = cuda_kernel_desc{axpy, 16, 128, 0, alpha, dX, dY}.get_num_registers();
+
+    int nregs_fun = cuda_kernel_desc{axpy_fun, 16, 128, 0, alpha, dX, dY}.get_num_registers();
+    _CCCL_ASSERT(nregs == nregs_fun, "invalid value");
+
+#if _CCCL_CTK_AT_LEAST(12, 1)
+    int nregs_kernel = cuda_kernel_desc{axpy_kernel, 16, 128, 0, alpha, dX, dY}.get_num_registers();
+    _CCCL_ASSERT(nregs == nregs_kernel, "invalid value");
+#endif
+  };
+
+  ctx.finalize();
+
+  for (size_t i = 0; i < N; i++)
+  {
+    _CCCL_ASSERT(fabs(Y[i] - (Y0(i) + num_axpy * alpha * X0(i))) < 0.0001, "Invalid result");
+    _CCCL_ASSERT(fabs(X[i] - X0(i)) < 0.0001, "Invalid result");
+  }
+}
+
+int main()
+{
+  // stream context
+  test(false);
+  // graph context
+  test(true);
+}

cccl_upstream/cudax/test/stf/fhe/batcher_sort_interpreted_cufhe.bool.ir

@@ -0,0 +1,355 @@
+package my_package
+
+fn conditional_swap(x: bits[16], y: bits[16], tmp: bits[16]) -> bits[16] {
+  bit_slice.936: bits[1] = bit_slice(y, start=15, width=1, id=936)
+  bit_slice.935: bits[1] = bit_slice(y, start=14, width=1, id=935)
+  bit_slice.952: bits[1] = bit_slice(x, start=15, width=1, id=952)
+  not.953: bits[1] = not(bit_slice.936, id=953)
+  bit_slice.951: bits[1] = bit_slice(x, start=14, width=1, id=951)
+  not.954: bits[1] = not(bit_slice.935, id=954)
+  bit_slice.934: bits[1] = bit_slice(y, start=13, width=1, id=934)
+  and.968: bits[1] = and(bit_slice.952, not.953, id=968)
+  and.971: bits[1] = and(bit_slice.951, not.954, id=971)
+  bit_slice.950: bits[1] = bit_slice(x, start=13, width=1, id=950)
+  not.955: bits[1] = not(bit_slice.934, id=955)
+  bit_slice.933: bits[1] = bit_slice(y, start=12, width=1, id=933)
+  not.969: bits[1] = not(and.968, id=969)
+  not.970: bits[1] = not(bit_slice.951, id=970)
+  not.972: bits[1] = not(and.971, id=972)
+  and.974: bits[1] = and(bit_slice.950, not.955, id=974)
+  bit_slice.949: bits[1] = bit_slice(x, start=12, width=1, id=949)
+  not.956: bits[1] = not(bit_slice.933, id=956)
+  bit_slice.932: bits[1] = bit_slice(y, start=11, width=1, id=932)
+  or.1014: bits[1] = or(bit_slice.952, not.953, id=1014)
+  and.1016: bits[1] = and(not.969, not.970, id=1016)
+  and.1018: bits[1] = and(not.969, not.972, id=1018)
+  not.973: bits[1] = not(bit_slice.950, id=973)
+  and.1021: bits[1] = and(not.969, not.972, id=1021)
+  not.975: bits[1] = not(and.974, id=975)
+  and.1025: bits[1] = and(not.969, not.972, id=1025)
+  and.977: bits[1] = and(bit_slice.949, not.956, id=977)
+  and.1030: bits[1] = and(not.969, not.972, id=1030)
+  bit_slice.948: bits[1] = bit_slice(x, start=11, width=1, id=948)
+  not.957: bits[1] = not(bit_slice.932, id=957)
+  and.1036: bits[1] = and(not.969, not.972, id=1036)
+  bit_slice.931: bits[1] = bit_slice(y, start=10, width=1, id=931)
+  and.1043: bits[1] = and(not.969, not.972, id=1043)
+  and.1051: bits[1] = and(not.969, not.972, id=1051)
+  and.1060: bits[1] = and(not.969, not.972, id=1060)
+  and.1070: bits[1] = and(not.969, not.972, id=1070)
+  and.1081: bits[1] = and(not.969, not.972, id=1081)
+  and.1093: bits[1] = and(not.969, not.972, id=1093)
+  and.1106: bits[1] = and(not.969, not.972, id=1106)
+  and.1120: bits[1] = and(not.969, not.972, id=1120)
+  and.1135: bits[1] = and(not.969, not.972, id=1135)
+  not.1015: bits[1] = not(or.1014, id=1015)
+  and.1017: bits[1] = and(and.1016, bit_slice.935, id=1017)
+  and.1019: bits[1] = and(and.1018, not.973, id=1019)
+  and.1022: bits[1] = and(and.1021, not.975, id=1022)
+  not.976: bits[1] = not(bit_slice.949, id=976)
+  and.1026: bits[1] = and(and.1025, not.975, id=1026)
+  not.978: bits[1] = not(and.977, id=978)
+  and.1031: bits[1] = and(and.1030, not.975, id=1031)
+  and.980: bits[1] = and(bit_slice.948, not.957, id=980)
+  and.1037: bits[1] = and(and.1036, not.975, id=1037)
+  bit_slice.947: bits[1] = bit_slice(x, start=10, width=1, id=947)
+  not.958: bits[1] = not(bit_slice.931, id=958)
+  and.1044: bits[1] = and(and.1043, not.975, id=1044)
+  bit_slice.930: bits[1] = bit_slice(y, start=9, width=1, id=930)
+  and.1052: bits[1] = and(and.1051, not.975, id=1052)
+  and.1061: bits[1] = and(and.1060, not.975, id=1061)
+  and.1071: bits[1] = and(and.1070, not.975, id=1071)
+  and.1082: bits[1] = and(and.1081, not.975, id=1082)
+  and.1094: bits[1] = and(and.1093, not.975, id=1094)
+  and.1107: bits[1] = and(and.1106, not.975, id=1107)
+  and.1121: bits[1] = and(and.1120, not.975, id=1121)
+  and.1136: bits[1] = and(and.1135, not.975, id=1136)
+  or.1153: bits[1] = or(not.1015, and.1017, id=1153)
+  and.1020: bits[1] = and(and.1019, bit_slice.934, id=1020)
+  and.1023: bits[1] = and(and.1022, not.976, id=1023)
+  and.1027: bits[1] = and(and.1026, not.978, id=1027)
+  not.979: bits[1] = not(bit_slice.948, id=979)
+  and.1032: bits[1] = and(and.1031, not.978, id=1032)
+  not.981: bits[1] = not(and.980, id=981)
+  and.1038: bits[1] = and(and.1037, not.978, id=1038)
+  and.983: bits[1] = and(bit_slice.947, not.958, id=983)
+  and.1045: bits[1] = and(and.1044, not.978, id=1045)
+  bit_slice.946: bits[1] = bit_slice(x, start=9, width=1, id=946)
+  not.959: bits[1] = not(bit_slice.930, id=959)
+  and.1053: bits[1] = and(and.1052, not.978, id=1053)
+  bit_slice.929: bits[1] = bit_slice(y, start=8, width=1, id=929)
+  and.1062: bits[1] = and(and.1061, not.978, id=1062)
+  and.1072: bits[1] = and(and.1071, not.978, id=1072)
+  and.1083: bits[1] = and(and.1082, not.978, id=1083)
+  and.1095: bits[1] = and(and.1094, not.978, id=1095)
+  and.1108: bits[1] = and(and.1107, not.978, id=1108)
+  and.1122: bits[1] = and(and.1121, not.978, id=1122)
+  and.1137: bits[1] = and(and.1136, not.978, id=1137)
+  or.1154: bits[1] = or(or.1153, and.1020, id=1154)
+  and.1024: bits[1] = and(and.1023, bit_slice.933, id=1024)
+  and.1028: bits[1] = and(and.1027, not.979, id=1028)
+  and.1033: bits[1] = and(and.1032, not.981, id=1033)
+  not.982: bits[1] = not(bit_slice.947, id=982)
+  and.1039: bits[1] = and(and.1038, not.981, id=1039)
+  not.984: bits[1] = not(and.983, id=984)
+  and.1046: bits[1] = and(and.1045, not.981, id=1046)
+  and.986: bits[1] = and(bit_slice.946, not.959, id=986)
+  and.1054: bits[1] = and(and.1053, not.981, id=1054)
+  bit_slice.945: bits[1] = bit_slice(x, start=8, width=1, id=945)
+  not.960: bits[1] = not(bit_slice.929, id=960)
+  and.1063: bits[1] = and(and.1062, not.981, id=1063)
+  bit_slice.928: bits[1] = bit_slice(y, start=7, width=1, id=928)
+  and.1073: bits[1] = and(and.1072, not.981, id=1073)
+  and.1084: bits[1] = and(and.1083, not.981, id=1084)
+  and.1096: bits[1] = and(and.1095, not.981, id=1096)
+  and.1109: bits[1] = and(and.1108, not.981, id=1109)
+  and.1123: bits[1] = and(and.1122, not.981, id=1123)
+  and.1138: bits[1] = and(and.1137, not.981, id=1138)
+  or.1155: bits[1] = or(or.1154, and.1024, id=1155)
+  and.1029: bits[1] = and(and.1028, bit_slice.932, id=1029)
+  and.1034: bits[1] = and(and.1033, not.982, id=1034)
+  and.1040: bits[1] = and(and.1039, not.984, id=1040)
+  not.985: bits[1] = not(bit_slice.946, id=985)
+  and.1047: bits[1] = and(and.1046, not.984, id=1047)
+  not.987: bits[1] = not(and.986, id=987)
+  and.1055: bits[1] = and(and.1054, not.984, id=1055)
+  and.989: bits[1] = and(bit_slice.945, not.960, id=989)
+  and.1064: bits[1] = and(and.1063, not.984, id=1064)
+  bit_slice.944: bits[1] = bit_slice(x, start=7, width=1, id=944)
+  not.961: bits[1] = not(bit_slice.928, id=961)
+  and.1074: bits[1] = and(and.1073, not.984, id=1074)
+  bit_slice.927: bits[1] = bit_slice(y, start=6, width=1, id=927)
+  and.1085: bits[1] = and(and.1084, not.984, id=1085)
+  and.1097: bits[1] = and(and.1096, not.984, id=1097)
+  and.1110: bits[1] = and(and.1109, not.984, id=1110)
+  and.1124: bits[1] = and(and.1123, not.984, id=1124)
+  and.1139: bits[1] = and(and.1138, not.984, id=1139)
+  or.1156: bits[1] = or(or.1155, and.1029, id=1156)
+  and.1035: bits[1] = and(and.1034, bit_slice.931, id=1035)
+  and.1041: bits[1] = and(and.1040, not.985, id=1041)
+  and.1048: bits[1] = and(and.1047, not.987, id=1048)
+  not.988: bits[1] = not(bit_slice.945, id=988)
+  and.1056: bits[1] = and(and.1055, not.987, id=1056)
+  not.990: bits[1] = not(and.989, id=990)
+  and.1065: bits[1] = and(and.1064, not.987, id=1065)
+  and.992: bits[1] = and(bit_slice.944, not.961, id=992)
+  and.1075: bits[1] = and(and.1074, not.987, id=1075)
+  bit_slice.943: bits[1] = bit_slice(x, start=6, width=1, id=943)
+  not.962: bits[1] = not(bit_slice.927, id=962)
+  and.1086: bits[1] = and(and.1085, not.987, id=1086)
+  bit_slice.926: bits[1] = bit_slice(y, start=5, width=1, id=926)
+  and.1098: bits[1] = and(and.1097, not.987, id=1098)
+  and.1111: bits[1] = and(and.1110, not.987, id=1111)
+  and.1125: bits[1] = and(and.1124, not.987, id=1125)
+  and.1140: bits[1] = and(and.1139, not.987, id=1140)
+  or.1157: bits[1] = or(or.1156, and.1035, id=1157)
+  and.1042: bits[1] = and(and.1041, bit_slice.930, id=1042)
+  and.1049: bits[1] = and(and.1048, not.988, id=1049)
+  and.1057: bits[1] = and(and.1056, not.990, id=1057)
+  not.991: bits[1] = not(bit_slice.944, id=991)
+  and.1066: bits[1] = and(and.1065, not.990, id=1066)
+  not.993: bits[1] = not(and.992, id=993)
+  and.1076: bits[1] = and(and.1075, not.990, id=1076)
+  and.995: bits[1] = and(bit_slice.943, not.962, id=995)
+  and.1087: bits[1] = and(and.1086, not.990, id=1087)
+  bit_slice.942: bits[1] = bit_slice(x, start=5, width=1, id=942)
+  not.963: bits[1] = not(bit_slice.926, id=963)
+  and.1099: bits[1] = and(and.1098, not.990, id=1099)
+  bit_slice.925: bits[1] = bit_slice(y, start=4, width=1, id=925)
+  and.1112: bits[1] = and(and.1111, not.990, id=1112)
+  and.1126: bits[1] = and(and.1125, not.990, id=1126)
+  and.1141: bits[1] = and(and.1140, not.990, id=1141)
+  or.1158: bits[1] = or(or.1157, and.1042, id=1158)
+  and.1050: bits[1] = and(and.1049, bit_slice.929, id=1050)
+  and.1058: bits[1] = and(and.1057, not.991, id=1058)
+  and.1067: bits[1] = and(and.1066, not.993, id=1067)
+  not.994: bits[1] = not(bit_slice.943, id=994)
+  and.1077: bits[1] = and(and.1076, not.993, id=1077)
+  not.996: bits[1] = not(and.995, id=996)
+  and.1088: bits[1] = and(and.1087, not.993, id=1088)
+  and.998: bits[1] = and(bit_slice.942, not.963, id=998)
+  and.1100: bits[1] = and(and.1099, not.993, id=1100)
+  bit_slice.941: bits[1] = bit_slice(x, start=4, width=1, id=941)
+  not.964: bits[1] = not(bit_slice.925, id=964)
+  and.1113: bits[1] = and(and.1112, not.993, id=1113)
+  bit_slice.924: bits[1] = bit_slice(y, start=3, width=1, id=924)
+  and.1127: bits[1] = and(and.1126, not.993, id=1127)
+  and.1142: bits[1] = and(and.1141, not.993, id=1142)
+  or.1159: bits[1] = or(or.1158, and.1050, id=1159)
+  and.1059: bits[1] = and(and.1058, bit_slice.928, id=1059)
+  and.1068: bits[1] = and(and.1067, not.994, id=1068)
+  and.1078: bits[1] = and(and.1077, not.996, id=1078)
+  not.997: bits[1] = not(bit_slice.942, id=997)
+  and.1089: bits[1] = and(and.1088, not.996, id=1089)
+  not.999: bits[1] = not(and.998, id=999)
+  and.1101: bits[1] = and(and.1100, not.996, id=1101)
+  and.1001: bits[1] = and(bit_slice.941, not.964, id=1001)
+  and.1114: bits[1] = and(and.1113, not.996, id=1114)
+  bit_slice.940: bits[1] = bit_slice(x, start=3, width=1, id=940)
+  not.965: bits[1] = not(bit_slice.924, id=965)
+  and.1128: bits[1] = and(and.1127, not.996, id=1128)
+  bit_slice.923: bits[1] = bit_slice(y, start=2, width=1, id=923)
+  and.1143: bits[1] = and(and.1142, not.996, id=1143)
+  or.1160: bits[1] = or(or.1159, and.1059, id=1160)
+  and.1069: bits[1] = and(and.1068, bit_slice.927, id=1069)
+  and.1079: bits[1] = and(and.1078, not.997, id=1079)
+  and.1090: bits[1] = and(and.1089, not.999, id=1090)
+  not.1000: bits[1] = not(bit_slice.941, id=1000)
+  and.1102: bits[1] = and(and.1101, not.999, id=1102)
+  not.1002: bits[1] = not(and.1001, id=1002)
+  and.1115: bits[1] = and(and.1114, not.999, id=1115)
+  and.1004: bits[1] = and(bit_slice.940, not.965, id=1004)
+  and.1129: bits[1] = and(and.1128, not.999, id=1129)
+  bit_slice.939: bits[1] = bit_slice(x, start=2, width=1, id=939)
+  not.966: bits[1] = not(bit_slice.923, id=966)
+  and.1144: bits[1] = and(and.1143, not.999, id=1144)
+  bit_slice.922: bits[1] = bit_slice(y, start=1, width=1, id=922)
+  or.1161: bits[1] = or(or.1160, and.1069, id=1161)
+  and.1080: bits[1] = and(and.1079, bit_slice.926, id=1080)
+  and.1091: bits[1] = and(and.1090, not.1000, id=1091)
+  and.1103: bits[1] = and(and.1102, not.1002, id=1103)
+  not.1003: bits[1] = not(bit_slice.940, id=1003)
+  and.1116: bits[1] = and(and.1115, not.1002, id=1116)
+  not.1005: bits[1] = not(and.1004, id=1005)
+  and.1130: bits[1] = and(and.1129, not.1002, id=1130)
+  and.1007: bits[1] = and(bit_slice.939, not.966, id=1007)
+  and.1145: bits[1] = and(and.1144, not.1002, id=1145)
+  bit_slice.938: bits[1] = bit_slice(x, start=1, width=1, id=938)
+  not.967: bits[1] = not(bit_slice.922, id=967)
+  or.1162: bits[1] = or(or.1161, and.1080, id=1162)
+  and.1092: bits[1] = and(and.1091, bit_slice.925, id=1092)
+  and.1104: bits[1] = and(and.1103, not.1003, id=1104)
+  and.1117: bits[1] = and(and.1116, not.1005, id=1117)
+  not.1006: bits[1] = not(bit_slice.939, id=1006)
+  and.1131: bits[1] = and(and.1130, not.1005, id=1131)
+  not.1008: bits[1] = not(and.1007, id=1008)
+  and.1146: bits[1] = and(and.1145, not.1005, id=1146)
+  and.1010: bits[1] = and(bit_slice.938, not.967, id=1010)
+  or.1163: bits[1] = or(or.1162, and.1092, id=1163)
+  and.1105: bits[1] = and(and.1104, bit_slice.924, id=1105)
+  and.1118: bits[1] = and(and.1117, not.1006, id=1118)
+  and.1132: bits[1] = and(and.1131, not.1008, id=1132)
+  not.1009: bits[1] = not(bit_slice.938, id=1009)
+  and.1147: bits[1] = and(and.1146, not.1008, id=1147)
+  not.1011: bits[1] = not(and.1010, id=1011)
+  bit_slice.937: bits[1] = bit_slice(x, start=0, width=1, id=937)
+  or.1164: bits[1] = or(or.1163, and.1105, id=1164)
+  and.1119: bits[1] = and(and.1118, bit_slice.923, id=1119)
+  and.1133: bits[1] = and(and.1132, not.1009, id=1133)
+  and.1148: bits[1] = and(and.1147, not.1011, id=1148)
+  not.1012: bits[1] = not(bit_slice.937, id=1012)
+  not.1013: bits[1] = not(bit_slice.952, id=1013)
+  or.1165: bits[1] = or(or.1164, and.1119, id=1165)
+  and.1134: bits[1] = and(and.1133, bit_slice.922, id=1134)
+  and.1149: bits[1] = and(and.1148, not.1012, id=1149)
+  bit_slice.921: bits[1] = bit_slice(y, start=0, width=1, id=921)
+  and.1151: bits[1] = and(not.1013, bit_slice.936, id=1151)
+  or.1166: bits[1] = or(or.1165, and.1134, id=1166)
+  and.1150: bits[1] = and(and.1149, bit_slice.921, id=1150)
+  or.1168: bits[1] = or(not.1013, bit_slice.936, id=1168)
+  not.1152: bits[1] = not(and.1151, id=1152)
+  or.1167: bits[1] = or(or.1166, and.1150, id=1167)
+  not.1169: bits[1] = not(or.1168, id=1169)
+  and.1170: bits[1] = and(not.1152, or.1167, id=1170)
+  bit_slice.1186: bits[1] = bit_slice(tmp, start=15, width=1, id=1186)
+  or.1187: bits[1] = or(not.1169, and.1170, id=1187)
+  bit_slice.1185: bits[1] = bit_slice(tmp, start=14, width=1, id=1185)
+  bit_slice.1184: bits[1] = bit_slice(tmp, start=13, width=1, id=1184)
+  bit_slice.1183: bits[1] = bit_slice(tmp, start=12, width=1, id=1183)
+  bit_slice.1182: bits[1] = bit_slice(tmp, start=11, width=1, id=1182)
+  bit_slice.1181: bits[1] = bit_slice(tmp, start=10, width=1, id=1181)
+  bit_slice.1180: bits[1] = bit_slice(tmp, start=9, width=1, id=1180)
+  bit_slice.1179: bits[1] = bit_slice(tmp, start=8, width=1, id=1179)
+  bit_slice.1178: bits[1] = bit_slice(tmp, start=7, width=1, id=1178)
+  bit_slice.1177: bits[1] = bit_slice(tmp, start=6, width=1, id=1177)
+  bit_slice.1176: bits[1] = bit_slice(tmp, start=5, width=1, id=1176)
+  bit_slice.1175: bits[1] = bit_slice(tmp, start=4, width=1, id=1175)
+  bit_slice.1174: bits[1] = bit_slice(tmp, start=3, width=1, id=1174)
+  bit_slice.1173: bits[1] = bit_slice(tmp, start=2, width=1, id=1173)
+  bit_slice.1172: bits[1] = bit_slice(tmp, start=1, width=1, id=1172)
+  bit_slice.1171: bits[1] = bit_slice(tmp, start=0, width=1, id=1171)
+  and.1189: bits[1] = and(bit_slice.952, bit_slice.1186, id=1189)
+  and.1190: bits[1] = and(bit_slice.952, or.1187, id=1190)
+  not.1188: bits[1] = not(or.1187, id=1188)
+  and.1192: bits[1] = and(bit_slice.951, bit_slice.1185, id=1192)
+  and.1193: bits[1] = and(bit_slice.951, or.1187, id=1193)
+  and.1195: bits[1] = and(bit_slice.950, bit_slice.1184, id=1195)
+  and.1196: bits[1] = and(bit_slice.950, or.1187, id=1196)
+  and.1198: bits[1] = and(bit_slice.949, bit_slice.1183, id=1198)
+  and.1199: bits[1] = and(bit_slice.949, or.1187, id=1199)
+  and.1201: bits[1] = and(bit_slice.948, bit_slice.1182, id=1201)
+  and.1202: bits[1] = and(bit_slice.948, or.1187, id=1202)
+  and.1204: bits[1] = and(bit_slice.947, bit_slice.1181, id=1204)
+  and.1205: bits[1] = and(bit_slice.947, or.1187, id=1205)
+  and.1207: bits[1] = and(bit_slice.946, bit_slice.1180, id=1207)
+  and.1208: bits[1] = and(bit_slice.946, or.1187, id=1208)
+  and.1210: bits[1] = and(bit_slice.945, bit_slice.1179, id=1210)
+  and.1211: bits[1] = and(bit_slice.945, or.1187, id=1211)
+  and.1213: bits[1] = and(bit_slice.944, bit_slice.1178, id=1213)
+  and.1214: bits[1] = and(bit_slice.944, or.1187, id=1214)
+  and.1216: bits[1] = and(bit_slice.943, bit_slice.1177, id=1216)
+  and.1217: bits[1] = and(bit_slice.943, or.1187, id=1217)
+  and.1219: bits[1] = and(bit_slice.942, bit_slice.1176, id=1219)
+  and.1220: bits[1] = and(bit_slice.942, or.1187, id=1220)
+  and.1222: bits[1] = and(bit_slice.941, bit_slice.1175, id=1222)
+  and.1223: bits[1] = and(bit_slice.941, or.1187, id=1223)
+  and.1225: bits[1] = and(bit_slice.940, bit_slice.1174, id=1225)
+  and.1226: bits[1] = and(bit_slice.940, or.1187, id=1226)
+  and.1228: bits[1] = and(bit_slice.939, bit_slice.1173, id=1228)
+  and.1229: bits[1] = and(bit_slice.939, or.1187, id=1229)
+  and.1231: bits[1] = and(bit_slice.938, bit_slice.1172, id=1231)
+  and.1232: bits[1] = and(bit_slice.938, or.1187, id=1232)
+  and.1234: bits[1] = and(bit_slice.937, bit_slice.1171, id=1234)
+  and.1235: bits[1] = and(bit_slice.937, or.1187, id=1235)
+  or.1237: bits[1] = or(and.1189, and.1190, id=1237)
+  and.1191: bits[1] = and(bit_slice.1186, not.1188, id=1191)
+  or.1239: bits[1] = or(and.1192, and.1193, id=1239)
+  and.1194: bits[1] = and(bit_slice.1185, not.1188, id=1194)
+  or.1241: bits[1] = or(and.1195, and.1196, id=1241)
+  and.1197: bits[1] = and(bit_slice.1184, not.1188, id=1197)
+  or.1243: bits[1] = or(and.1198, and.1199, id=1243)
+  and.1200: bits[1] = and(bit_slice.1183, not.1188, id=1200)
+  or.1245: bits[1] = or(and.1201, and.1202, id=1245)
+  and.1203: bits[1] = and(bit_slice.1182, not.1188, id=1203)
+  or.1247: bits[1] = or(and.1204, and.1205, id=1247)
+  and.1206: bits[1] = and(bit_slice.1181, not.1188, id=1206)
+  or.1249: bits[1] = or(and.1207, and.1208, id=1249)
+  and.1209: bits[1] = and(bit_slice.1180, not.1188, id=1209)
+  or.1251: bits[1] = or(and.1210, and.1211, id=1251)
+  and.1212: bits[1] = and(bit_slice.1179, not.1188, id=1212)
+  or.1253: bits[1] = or(and.1213, and.1214, id=1253)
+  and.1215: bits[1] = and(bit_slice.1178, not.1188, id=1215)
+  or.1255: bits[1] = or(and.1216, and.1217, id=1255)
+  and.1218: bits[1] = and(bit_slice.1177, not.1188, id=1218)
+  or.1257: bits[1] = or(and.1219, and.1220, id=1257)
+  and.1221: bits[1] = and(bit_slice.1176, not.1188, id=1221)
+  or.1259: bits[1] = or(and.1222, and.1223, id=1259)
+  and.1224: bits[1] = and(bit_slice.1175, not.1188, id=1224)
+  or.1261: bits[1] = or(and.1225, and.1226, id=1261)
+  and.1227: bits[1] = and(bit_slice.1174, not.1188, id=1227)
+  or.1263: bits[1] = or(and.1228, and.1229, id=1263)
+  and.1230: bits[1] = and(bit_slice.1173, not.1188, id=1230)
+  or.1265: bits[1] = or(and.1231, and.1232, id=1265)
+  and.1233: bits[1] = and(bit_slice.1172, not.1188, id=1233)
+  or.1267: bits[1] = or(and.1234, and.1235, id=1267)
+  and.1236: bits[1] = and(bit_slice.1171, not.1188, id=1236)
+  or.1238: bits[1] = or(or.1237, and.1191, id=1238)
+  or.1240: bits[1] = or(or.1239, and.1194, id=1240)
+  or.1242: bits[1] = or(or.1241, and.1197, id=1242)
+  or.1244: bits[1] = or(or.1243, and.1200, id=1244)
+  or.1246: bits[1] = or(or.1245, and.1203, id=1246)
+  or.1248: bits[1] = or(or.1247, and.1206, id=1248)
+  or.1250: bits[1] = or(or.1249, and.1209, id=1250)
+  or.1252: bits[1] = or(or.1251, and.1212, id=1252)
+  or.1254: bits[1] = or(or.1253, and.1215, id=1254)
+  or.1256: bits[1] = or(or.1255, and.1218, id=1256)
+  or.1258: bits[1] = or(or.1257, and.1221, id=1258)
+  or.1260: bits[1] = or(or.1259, and.1224, id=1260)
+  or.1262: bits[1] = or(or.1261, and.1227, id=1262)
+  or.1264: bits[1] = or(or.1263, and.1230, id=1264)
+  or.1266: bits[1] = or(or.1265, and.1233, id=1266)
+  or.1268: bits[1] = or(or.1267, and.1236, id=1268)
+  literal.916: bits[1] = literal(value=1, id=916)
+  literal.917: bits[1] = literal(value=0, id=917)
+  ret concat.1269: bits[16] = concat(or.1238, or.1240, or.1242, or.1244, or.1246, or.1248, or.1250, or.1252, or.1254, or.1256, or.1258, or.1260, or.1262, or.1264, or.1266, or.1268, id=1269)
+}

cccl_upstream/cudax/test/stf/fhe/kernel_sharpen_interpreted_cufhe.bool.ir

@@ -0,0 +1,556 @@
+package my_package
+
+fn kernel_sharpen(window: bits[8][9]) -> bits[8] {
+  literal.2853: bits[4] = literal(value=5, id=2853)
+  literal.2873: bits[4] = literal(value=7, id=2873)
+  array_index.2854: bits[8] = array_index(window, indices=[literal.2853], id=2854)
+  array_index.2874: bits[8] = array_index(window, indices=[literal.2873], id=2874)
+  literal.2813: bits[4] = literal(value=1, id=2813)
+  literal.2833: bits[4] = literal(value=3, id=2833)
+  bit_slice.2857: bits[1] = bit_slice(array_index.2854, start=2, width=1, id=2857)
+  bit_slice.2877: bits[1] = bit_slice(array_index.2874, start=2, width=1, id=2877)
+  array_index.2814: bits[8] = array_index(window, indices=[literal.2813], id=2814)
+  array_index.2834: bits[8] = array_index(window, indices=[literal.2833], id=2834)
+  and.2894: bits[1] = and(bit_slice.2857, bit_slice.2877, id=2894)
+  bit_slice.2817: bits[1] = bit_slice(array_index.2814, start=2, width=1, id=2817)
+  bit_slice.2837: bits[1] = bit_slice(array_index.2834, start=2, width=1, id=2837)
+  or.2893: bits[1] = or(bit_slice.2857, bit_slice.2877, id=2893)
+  not.2895: bits[1] = not(and.2894, id=2895)
+  and.2897: bits[1] = and(bit_slice.2817, bit_slice.2837, id=2897)
+  and.2907: bits[1] = and(bit_slice.2857, bit_slice.2877, id=2907)
+  and.2900: bits[1] = and(or.2893, not.2895, id=2900)
+  bit_slice.2856: bits[1] = bit_slice(array_index.2854, start=1, width=1, id=2856)
+  or.2896: bits[1] = or(bit_slice.2817, bit_slice.2837, id=2896)
+  not.2898: bits[1] = not(and.2897, id=2898)
+  or.2909: bits[1] = or(bit_slice.2857, bit_slice.2877, id=2909)
+  not.2908: bits[1] = not(and.2907, id=2908)
+  bit_slice.2876: bits[1] = bit_slice(array_index.2874, start=1, width=1, id=2876)
+  and.2901: bits[1] = and(and.2900, bit_slice.2856, id=2901)
+  and.2912: bits[1] = and(bit_slice.2817, bit_slice.2837, id=2912)
+  and.2904: bits[1] = and(or.2896, not.2898, id=2904)
+  bit_slice.2816: bits[1] = bit_slice(array_index.2814, start=1, width=1, id=2816)
+  and.2910: bits[1] = and(or.2909, not.2908, id=2910)
+  and.2911: bits[1] = and(bit_slice.2856, bit_slice.2876, id=2911)
+  and.2902: bits[1] = and(and.2901, bit_slice.2876, id=2902)
+  or.2914: bits[1] = or(bit_slice.2817, bit_slice.2837, id=2914)
+  not.2913: bits[1] = not(and.2912, id=2913)
+  bit_slice.2836: bits[1] = bit_slice(array_index.2834, start=1, width=1, id=2836)
+  bit_slice.2858: bits[1] = bit_slice(array_index.2854, start=3, width=1, id=2858)
+  bit_slice.2878: bits[1] = bit_slice(array_index.2874, start=3, width=1, id=2878)
+  and.2905: bits[1] = and(and.2904, bit_slice.2816, id=2905)
+  or.2925: bits[1] = or(and.2910, and.2911, id=2925)
+  not.2926: bits[1] = not(and.2902, id=2926)
+  and.2915: bits[1] = and(or.2914, not.2913, id=2915)
+  and.2916: bits[1] = and(bit_slice.2816, bit_slice.2836, id=2916)
+  and.2918: bits[1] = and(bit_slice.2858, bit_slice.2878, id=2918)
+  and.2935: bits[1] = and(bit_slice.2858, bit_slice.2878, id=2935)
+  bit_slice.2818: bits[1] = bit_slice(array_index.2814, start=3, width=1, id=2818)
+  bit_slice.2838: bits[1] = bit_slice(array_index.2834, start=3, width=1, id=2838)
+  and.2906: bits[1] = and(and.2905, bit_slice.2836, id=2906)
+  and.2945: bits[1] = and(or.2925, not.2926, id=2945)
+  or.2927: bits[1] = or(and.2915, and.2916, id=2927)
+  bit_slice.2859: bits[1] = bit_slice(array_index.2854, start=4, width=1, id=2859)
+  bit_slice.2879: bits[1] = bit_slice(array_index.2874, start=4, width=1, id=2879)
+  or.2917: bits[1] = or(bit_slice.2858, bit_slice.2878, id=2917)
+  not.2919: bits[1] = not(and.2918, id=2919)
+  and.2899: bits[1] = and(bit_slice.2857, bit_slice.2877, id=2899)
+  or.2937: bits[1] = or(bit_slice.2858, bit_slice.2878, id=2937)
+  not.2936: bits[1] = not(and.2935, id=2936)
+  and.2939: bits[1] = and(bit_slice.2818, bit_slice.2838, id=2939)
+  not.2928: bits[1] = not(and.2906, id=2928)
+  and.2946: bits[1] = and(and.2945, or.2927, id=2946)
+  and.2922: bits[1] = and(bit_slice.2818, bit_slice.2838, id=2922)
+  and.2949: bits[1] = and(bit_slice.2859, bit_slice.2879, id=2949)
+  and.2930: bits[1] = and(or.2917, not.2919, id=2930)
+  or.2920: bits[1] = or(and.2899, and.2902, id=2920)
+  and.2976: bits[1] = and(bit_slice.2859, bit_slice.2879, id=2976)
+  bit_slice.2819: bits[1] = bit_slice(array_index.2814, start=4, width=1, id=2819)
+  bit_slice.2839: bits[1] = bit_slice(array_index.2834, start=4, width=1, id=2839)
+  and.2938: bits[1] = and(or.2937, not.2936, id=2938)
+  or.2941: bits[1] = or(bit_slice.2818, bit_slice.2838, id=2941)
+  not.2940: bits[1] = not(and.2939, id=2940)
+  and.2943: bits[1] = and(or.2925, not.2926, id=2943)
+  and.2944: bits[1] = and(or.2927, not.2928, id=2944)
+  and.2947: bits[1] = and(and.2946, not.2928, id=2947)
+  or.2921: bits[1] = or(bit_slice.2818, bit_slice.2838, id=2921)
+  not.2923: bits[1] = not(and.2922, id=2923)
+  and.2903: bits[1] = and(bit_slice.2817, bit_slice.2837, id=2903)
+  bit_slice.2860: bits[1] = bit_slice(array_index.2854, start=5, width=1, id=2860)
+  bit_slice.2880: bits[1] = bit_slice(array_index.2874, start=5, width=1, id=2880)
+  or.2948: bits[1] = or(bit_slice.2859, bit_slice.2879, id=2948)
+  not.2950: bits[1] = not(and.2949, id=2950)
+  and.2929: bits[1] = and(bit_slice.2858, bit_slice.2878, id=2929)
+  and.2931: bits[1] = and(and.2930, or.2920, id=2931)
+  or.2978: bits[1] = or(bit_slice.2859, bit_slice.2879, id=2978)
+  not.2977: bits[1] = not(and.2976, id=2977)
+  and.2980: bits[1] = and(bit_slice.2819, bit_slice.2839, id=2980)
+  or.2956: bits[1] = or(and.2938, and.2899, id=2956)
+  and.2942: bits[1] = and(or.2941, not.2940, id=2942)
+  or.2962: bits[1] = or(and.2943, and.2944, id=2962)
+  not.2963: bits[1] = not(and.2947, id=2963)
+  and.2953: bits[1] = and(bit_slice.2819, bit_slice.2839, id=2953)
+  and.2933: bits[1] = and(or.2921, not.2923, id=2933)
+  or.2924: bits[1] = or(and.2903, and.2906, id=2924)
+  and.2995: bits[1] = and(bit_slice.2860, bit_slice.2880, id=2995)
+  and.2971: bits[1] = and(or.2948, not.2950, id=2971)
+  or.2951: bits[1] = or(and.2929, and.2931, id=2951)
+  and.3017: bits[1] = and(bit_slice.2860, bit_slice.2880, id=3017)
+  bit_slice.2820: bits[1] = bit_slice(array_index.2814, start=5, width=1, id=2820)
+  bit_slice.2840: bits[1] = bit_slice(array_index.2834, start=5, width=1, id=2840)
+  and.2979: bits[1] = and(or.2978, not.2977, id=2979)
+  or.2982: bits[1] = or(bit_slice.2819, bit_slice.2839, id=2982)
+  not.2981: bits[1] = not(and.2980, id=2981)
+  or.2957: bits[1] = or(or.2956, and.2902, id=2957)
+  not.2958: bits[1] = not(and.2931, id=2958)
+  or.2959: bits[1] = or(and.2942, and.2903, id=2959)
+  and.2989: bits[1] = and(or.2962, not.2963, id=2989)
+  or.2964: bits[1] = or(bit_slice.2856, bit_slice.2876, id=2964)
+  and.2965: bits[1] = and(bit_slice.2856, bit_slice.2876, id=2965)
+  or.2952: bits[1] = or(bit_slice.2819, bit_slice.2839, id=2952)
+  not.2954: bits[1] = not(and.2953, id=2954)
+  and.2932: bits[1] = and(bit_slice.2818, bit_slice.2838, id=2932)
+  and.2934: bits[1] = and(and.2933, or.2924, id=2934)
+  bit_slice.2861: bits[1] = bit_slice(array_index.2854, start=6, width=1, id=2861)
+  bit_slice.2881: bits[1] = bit_slice(array_index.2874, start=6, width=1, id=2881)
+  or.2994: bits[1] = or(bit_slice.2860, bit_slice.2880, id=2994)
+  not.2996: bits[1] = not(and.2995, id=2996)
+  and.2970: bits[1] = and(bit_slice.2859, bit_slice.2879, id=2970)
+  and.2972: bits[1] = and(and.2971, or.2951, id=2972)
+  or.3019: bits[1] = or(bit_slice.2860, bit_slice.2880, id=3019)
+  not.3018: bits[1] = not(and.3017, id=3018)
+  and.3021: bits[1] = and(bit_slice.2820, bit_slice.2840, id=3021)
+  or.3002: bits[1] = or(and.2979, and.2929, id=3002)
+  and.2983: bits[1] = and(or.2982, not.2981, id=2983)
+  and.2986: bits[1] = and(or.2957, not.2958, id=2986)
+  or.2960: bits[1] = or(or.2959, and.2906, id=2960)
+  and.2990: bits[1] = and(and.2989, or.2964, id=2990)
+  not.2966: bits[1] = not(and.2965, id=2966)
+  and.2999: bits[1] = and(bit_slice.2820, bit_slice.2840, id=2999)
+  and.2974: bits[1] = and(or.2952, not.2954, id=2974)
+  or.2955: bits[1] = or(and.2932, and.2934, id=2955)
+  and.3033: bits[1] = and(bit_slice.2861, bit_slice.2881, id=3033)
+  and.3012: bits[1] = and(or.2994, not.2996, id=3012)
+  or.2997: bits[1] = or(and.2970, and.2972, id=2997)
+  and.3055: bits[1] = and(bit_slice.2861, bit_slice.2881, id=3055)
+  bit_slice.2821: bits[1] = bit_slice(array_index.2814, start=6, width=1, id=2821)
+  bit_slice.2841: bits[1] = bit_slice(array_index.2834, start=6, width=1, id=2841)
+  and.3020: bits[1] = and(or.3019, not.3018, id=3020)
+  or.3023: bits[1] = or(bit_slice.2820, bit_slice.2840, id=3023)
+  not.3022: bits[1] = not(and.3021, id=3022)
+  or.3003: bits[1] = or(or.3002, and.2931, id=3003)
+  not.3004: bits[1] = not(and.2972, id=3004)
+  or.3005: bits[1] = or(and.2983, and.2932, id=3005)
+  not.2961: bits[1] = not(and.2934, id=2961)
+  and.2987: bits[1] = and(and.2986, or.2960, id=2987)
+  and.2991: bits[1] = and(and.2990, not.2966, id=2991)
+  or.2967: bits[1] = or(bit_slice.2816, bit_slice.2836, id=2967)
+  and.2968: bits[1] = and(bit_slice.2816, bit_slice.2836, id=2968)
+  or.2998: bits[1] = or(bit_slice.2820, bit_slice.2840, id=2998)
+  not.3000: bits[1] = not(and.2999, id=3000)
+  and.2973: bits[1] = and(bit_slice.2819, bit_slice.2839, id=2973)
+  and.2975: bits[1] = and(and.2974, or.2955, id=2975)
+  bit_slice.2862: bits[1] = bit_slice(array_index.2854, start=7, width=1, id=2862)
+  bit_slice.2882: bits[1] = bit_slice(array_index.2874, start=7, width=1, id=2882)
+  or.3032: bits[1] = or(bit_slice.2861, bit_slice.2881, id=3032)
+  not.3034: bits[1] = not(and.3033, id=3034)
+  and.3011: bits[1] = and(bit_slice.2860, bit_slice.2880, id=3011)
+  and.3013: bits[1] = and(and.3012, or.2997, id=3013)
+  or.3057: bits[1] = or(bit_slice.2861, bit_slice.2881, id=3057)
+  not.3056: bits[1] = not(and.3055, id=3056)
+  and.3059: bits[1] = and(bit_slice.2821, bit_slice.2841, id=3059)
+  or.3040: bits[1] = or(and.3020, and.2970, id=3040)
+  and.3024: bits[1] = and(or.3023, not.3022, id=3024)
+  and.3027: bits[1] = and(or.3003, not.3004, id=3027)
+  or.3006: bits[1] = or(or.3005, and.2934, id=3006)
+  and.2984: bits[1] = and(or.2957, not.2958, id=2984)
+  and.2985: bits[1] = and(or.2960, not.2961, id=2985)
+  and.2988: bits[1] = and(and.2987, not.2961, id=2988)
+  and.2992: bits[1] = and(and.2991, or.2967, id=2992)
+  not.2969: bits[1] = not(and.2968, id=2969)
+  and.3037: bits[1] = and(bit_slice.2821, bit_slice.2841, id=3037)
+  and.3015: bits[1] = and(or.2998, not.3000, id=3015)
+  or.3001: bits[1] = or(and.2973, and.2975, id=3001)
+  and.3087: bits[1] = and(bit_slice.2862, bit_slice.2882, id=3087)
+  and.3071: bits[1] = and(bit_slice.2862, bit_slice.2882, id=3071)
+  and.3050: bits[1] = and(or.3032, not.3034, id=3050)
+  or.3035: bits[1] = or(and.3011, and.3013, id=3035)
+  bit_slice.2822: bits[1] = bit_slice(array_index.2814, start=7, width=1, id=2822)
+  bit_slice.2842: bits[1] = bit_slice(array_index.2834, start=7, width=1, id=2842)
+  and.3058: bits[1] = and(or.3057, not.3056, id=3058)
+  or.3061: bits[1] = or(bit_slice.2821, bit_slice.2841, id=3061)
+  not.3060: bits[1] = not(and.3059, id=3060)
+  or.3041: bits[1] = or(or.3040, and.2972, id=3041)
+  not.3042: bits[1] = not(and.3013, id=3042)
+  or.3043: bits[1] = or(and.3024, and.2973, id=3043)
+  not.3007: bits[1] = not(and.2975, id=3007)
+  and.3028: bits[1] = and(and.3027, or.3006, id=3028)
+  or.3008: bits[1] = or(and.2984, and.2985, id=3008)
+  not.3009: bits[1] = not(and.2988, id=3009)
+  and.2993: bits[1] = and(and.2992, not.2969, id=2993)
+  or.3036: bits[1] = or(bit_slice.2821, bit_slice.2841, id=3036)
+  not.3038: bits[1] = not(and.3037, id=3038)
+  and.3014: bits[1] = and(bit_slice.2820, bit_slice.2840, id=3014)
+  and.3016: bits[1] = and(and.3015, or.3001, id=3016)
+  or.3089: bits[1] = or(bit_slice.2862, bit_slice.2882, id=3089)
+  not.3088: bits[1] = not(and.3087, id=3088)
+  or.3070: bits[1] = or(bit_slice.2862, bit_slice.2882, id=3070)
+  not.3072: bits[1] = not(and.3071, id=3072)
+  and.3049: bits[1] = and(bit_slice.2861, bit_slice.2881, id=3049)
+  and.3051: bits[1] = and(and.3050, or.3035, id=3051)
+  and.3093: bits[1] = and(bit_slice.2822, bit_slice.2842, id=3093)
+  or.3078: bits[1] = or(and.3058, and.3011, id=3078)
+  and.3062: bits[1] = and(or.3061, not.3060, id=3062)
+  and.3065: bits[1] = and(or.3041, not.3042, id=3065)
+  or.3044: bits[1] = or(or.3043, and.2975, id=3044)
+  and.3025: bits[1] = and(or.3003, not.3004, id=3025)
+  and.3026: bits[1] = and(or.3006, not.3007, id=3026)
+  and.3029: bits[1] = and(and.3028, not.3007, id=3029)
+  and.3030: bits[1] = and(or.3008, not.3009, id=3030)
+  or.3010: bits[1] = or(and.2947, and.2993, id=3010)
+  and.3075: bits[1] = and(bit_slice.2822, bit_slice.2842, id=3075)
+  and.3053: bits[1] = and(or.3036, not.3038, id=3053)
+  or.3039: bits[1] = or(and.3014, and.3016, id=3039)
+  and.3090: bits[1] = and(or.3089, not.3088, id=3090)
+  and.3091: bits[1] = and(or.3070, not.3072, id=3091)
+  or.3073: bits[1] = or(and.3049, and.3051, id=3073)
+  or.3095: bits[1] = or(bit_slice.2822, bit_slice.2842, id=3095)
+  not.3094: bits[1] = not(and.3093, id=3094)
+  or.3079: bits[1] = or(or.3078, and.3013, id=3079)
+  not.3080: bits[1] = not(and.3051, id=3080)
+  or.3081: bits[1] = or(and.3062, and.3014, id=3081)
+  not.3045: bits[1] = not(and.3016, id=3045)
+  and.3066: bits[1] = and(and.3065, or.3044, id=3066)
+  or.3046: bits[1] = or(and.3025, and.3026, id=3046)
+  not.3047: bits[1] = not(and.3029, id=3047)
+  and.3031: bits[1] = and(and.3030, or.3010, id=3031)
+  or.3074: bits[1] = or(bit_slice.2822, bit_slice.2842, id=3074)
+  not.3076: bits[1] = not(and.3075, id=3076)
+  and.3052: bits[1] = and(bit_slice.2821, bit_slice.2841, id=3052)
+  and.3054: bits[1] = and(and.3053, or.3039, id=3054)
+  or.3121: bits[1] = or(and.3090, and.3049, id=3121)
+  and.3092: bits[1] = and(and.3091, or.3073, id=3092)
+  and.3096: bits[1] = and(or.3095, not.3094, id=3096)
+  and.3101: bits[1] = and(or.3079, not.3080, id=3101)
+  or.3082: bits[1] = or(or.3081, and.3016, id=3082)
+  and.3063: bits[1] = and(or.3041, not.3042, id=3063)
+  and.3064: bits[1] = and(or.3044, not.3045, id=3064)
+  and.3067: bits[1] = and(and.3066, not.3045, id=3067)
+  and.3068: bits[1] = and(or.3046, not.3047, id=3068)
+  or.3048: bits[1] = or(and.2988, and.3031, id=3048)
+  and.3097: bits[1] = and(or.3074, not.3076, id=3097)
+  or.3077: bits[1] = or(and.3052, and.3054, id=3077)
+  or.3122: bits[1] = or(or.3121, and.3051, id=3122)
+  not.3123: bits[1] = not(and.3092, id=3123)
+  or.3124: bits[1] = or(and.3096, and.3052, id=3124)
+  not.3083: bits[1] = not(and.3054, id=3083)
+  and.3102: bits[1] = and(and.3101, or.3082, id=3102)
+  or.3084: bits[1] = or(and.3063, and.3064, id=3084)
+  not.3085: bits[1] = not(and.3067, id=3085)
+  and.3069: bits[1] = and(and.3068, or.3048, id=3069)
+  and.3098: bits[1] = and(and.3097, or.3077, id=3098)
+  and.3147: bits[1] = and(or.3122, not.3123, id=3147)
+  or.3125: bits[1] = or(or.3124, and.3054, id=3125)
+  and.3099: bits[1] = and(or.3079, not.3080, id=3099)
+  and.3100: bits[1] = and(or.3082, not.3083, id=3100)
+  and.3103: bits[1] = and(and.3102, not.3083, id=3103)
+  and.3104: bits[1] = and(or.3084, not.3085, id=3104)
+  or.3086: bits[1] = or(and.3029, and.3069, id=3086)
+  not.3126: bits[1] = not(and.3098, id=3126)
+  and.3148: bits[1] = and(and.3147, or.3125, id=3148)
+  or.3127: bits[1] = or(and.3099, and.3100, id=3127)
+  not.3128: bits[1] = not(and.3103, id=3128)
+  and.3105: bits[1] = and(and.3104, or.3086, id=3105)
+  and.3120: bits[1] = and(or.3084, not.3085, id=3120)
+  and.3106: bits[1] = and(or.3046, not.3047, id=3106)
+  and.3107: bits[1] = and(or.3008, not.3009, id=3107)
+  and.3109: bits[1] = and(or.2964, not.2966, id=3109)
+  and.3143: bits[1] = and(bit_slice.2862, bit_slice.2882, id=3143)
+  and.3144: bits[1] = and(bit_slice.2822, bit_slice.2842, id=3144)
+  and.3145: bits[1] = and(or.3122, not.3123, id=3145)
+  and.3146: bits[1] = and(or.3125, not.3126, id=3146)
+  and.3149: bits[1] = and(and.3148, not.3126, id=3149)
+  and.3150: bits[1] = and(or.3127, not.3128, id=3150)
+  or.3129: bits[1] = or(and.3067, and.3105, id=3129)
+  or.3140: bits[1] = or(and.3120, and.3029, id=3140)
+  or.3130: bits[1] = or(and.3106, and.2988, id=3130)
+  or.3133: bits[1] = or(and.3107, and.2947, id=3133)
+  and.3110: bits[1] = and(and.3109, or.2967, id=3110)
+  and.3112: bits[1] = and(bit_slice.2856, bit_slice.2876, id=3112)
+  and.3116: bits[1] = and(bit_slice.2816, bit_slice.2836, id=3116)
+  or.3183: bits[1] = or(and.3143, and.3092, id=3183)
+  or.3158: bits[1] = or(and.3143, and.3092, id=3158)
+  or.3159: bits[1] = or(and.3144, and.3098, id=3159)
+  or.3160: bits[1] = or(and.3145, and.3146, id=3160)
+  not.3161: bits[1] = not(and.3149, id=3161)
+  and.3151: bits[1] = and(and.3150, or.3129, id=3151)
+  or.3141: bits[1] = or(or.3140, and.3069, id=3141)
+  not.3142: bits[1] = not(and.3105, id=3142)
+  or.3131: bits[1] = or(or.3130, and.3031, id=3131)
+  not.3132: bits[1] = not(and.3069, id=3132)
+  or.3134: bits[1] = or(or.3133, and.2993, id=3134)
+  not.3135: bits[1] = not(and.3031, id=3135)
+  and.3108: bits[1] = and(or.2962, not.2963, id=3108)
+  and.3111: bits[1] = and(and.3110, not.2969, id=3111)
+  or.3114: bits[1] = or(bit_slice.2856, bit_slice.2876, id=3114)
+  not.3113: bits[1] = not(and.3112, id=3113)
+  or.3118: bits[1] = or(bit_slice.2816, bit_slice.2836, id=3118)
+  not.3117: bits[1] = not(and.3116, id=3117)
+  or.3184: bits[1] = or(or.3183, and.3144, id=3184)
+  and.3171: bits[1] = and(or.3158, or.3159, id=3171)
+  and.3172: bits[1] = and(or.3160, not.3161, id=3172)
+  or.3162: bits[1] = or(and.3103, and.3151, id=3162)
+  and.3152: bits[1] = and(or.3127, not.3128, id=3152)
+  and.3157: bits[1] = and(or.3141, not.3142, id=3157)
+  and.3153: bits[1] = and(or.3131, not.3132, id=3153)
+  and.3154: bits[1] = and(or.3134, not.3135, id=3154)
+  or.3136: bits[1] = or(and.3108, and.3111, id=3136)
+  not.3137: bits[1] = not(and.2993, id=3137)
+  and.3115: bits[1] = and(or.3114, not.3113, id=3115)
+  and.3119: bits[1] = and(or.3118, not.3117, id=3119)
+  or.3185: bits[1] = or(or.3184, and.3098, id=3185)
+  not.3186: bits[1] = not(and.3171, id=3186)
+  and.3173: bits[1] = and(and.3172, or.3162, id=3173)
+  and.3174: bits[1] = and(or.3160, not.3161, id=3174)
+  or.3163: bits[1] = or(and.3152, and.3067, id=3163)
+  not.3170: bits[1] = not(and.3157, id=3170)
+  not.3166: bits[1] = not(and.3153, id=3166)
+  not.3167: bits[1] = not(and.3154, id=3167)
+  and.3155: bits[1] = and(or.3136, not.3137, id=3155)
+  or.3138: bits[1] = or(and.3115, and.3119, id=3138)
+  not.3139: bits[1] = not(and.3111, id=3139)
+  and.3194: bits[1] = and(or.3185, not.3186, id=3194)
+  and.3195: bits[1] = and(or.3185, not.3186, id=3195)
+  or.3187: bits[1] = or(and.3149, and.3173, id=3187)
+  or.3188: bits[1] = or(and.3174, and.3103, id=3188)
+  or.3164: bits[1] = or(or.3163, and.3105, id=3164)
+  not.3165: bits[1] = not(and.3151, id=3165)
+  and.3179: bits[1] = and(not.3170, not.3166, id=3179)
+  and.3176: bits[1] = and(not.3166, not.3167, id=3176)
+  not.3168: bits[1] = not(and.3155, id=3168)
+  and.3156: bits[1] = and(or.3138, not.3139, id=3156)
+  or.3208: bits[1] = or(and.3194, and.3149, id=3208)
+  and.3196: bits[1] = and(and.3195, or.3187, id=3196)
+  or.3189: bits[1] = or(or.3188, and.3151, id=3189)
+  not.3190: bits[1] = not(and.3173, id=3190)
+  and.3175: bits[1] = and(or.3164, not.3165, id=3175)
+  and.3180: bits[1] = and(and.3179, not.3167, id=3180)
+  literal.2843: bits[4] = literal(value=4, id=2843)
+  and.3177: bits[1] = and(and.3176, not.3168, id=3177)
+  not.3169: bits[1] = not(and.3156, id=3169)
+  or.3209: bits[1] = or(or.3208, and.3173, id=3209)
+  not.3210: bits[1] = not(and.3196, id=3210)
+  and.3197: bits[1] = and(or.3189, not.3190, id=3197)
+  not.3191: bits[1] = not(and.3175, id=3191)
+  and.3181: bits[1] = and(and.3180, not.3168, id=3181)
+  array_index.2844: bits[8] = array_index(window, indices=[literal.2843], id=2844)
+  and.3178: bits[1] = and(and.3177, not.3169, id=3178)
+  and.3218: bits[1] = and(or.3209, not.3210, id=3218)
+  not.3211: bits[1] = not(and.3197, id=3211)
+  and.3198: bits[1] = and(not.3191, not.3170, id=3198)
+  and.3182: bits[1] = and(and.3181, not.3169, id=3182)
+  bit_slice.2846: bits[1] = bit_slice(array_index.2844, start=1, width=1, id=2846)
+  or.3192: bits[1] = or(not.3170, and.3178, id=3192)
+  not.3232: bits[1] = not(and.3218, id=3232)
+  and.3219: bits[1] = and(not.3211, not.3191, id=3219)
+  and.3199: bits[1] = and(and.3198, not.3166, id=3199)
+  not.3193: bits[1] = not(and.3182, id=3193)
+  and.3204: bits[1] = and(bit_slice.2846, or.3192, id=3204)
+  and.3238: bits[1] = and(not.3232, not.3211, id=3238)
+  and.3220: bits[1] = and(and.3219, not.3170, id=3220)
+  and.3200: bits[1] = and(and.3199, not.3167, id=3200)
+  and.3203: bits[1] = and(or.3192, not.3193, id=3203)
+  and.3205: bits[1] = and(and.3204, not.3193, id=3205)
+  and.3239: bits[1] = and(and.3238, not.3191, id=3239)
+  and.3221: bits[1] = and(and.3220, not.3166, id=3221)
+  and.3201: bits[1] = and(and.3200, not.3168, id=3201)
+  or.3214: bits[1] = or(bit_slice.2846, and.3203, id=3214)
+  not.3215: bits[1] = not(and.3205, id=3215)
+  and.3206: bits[1] = and(not.3167, not.3168, id=3206)
+  and.3240: bits[1] = and(and.3239, not.3170, id=3240)
+  and.3222: bits[1] = and(and.3221, not.3167, id=3222)
+  and.3202: bits[1] = and(and.3201, not.3169, id=3202)
+  bit_slice.2847: bits[1] = bit_slice(array_index.2844, start=2, width=1, id=2847)
+  or.3212: bits[1] = or(not.3191, and.3182, id=3212)
+  and.3228: bits[1] = and(or.3214, not.3215, id=3228)
+  bit_slice.2845: bits[1] = bit_slice(array_index.2844, start=0, width=1, id=2845)
+  and.3207: bits[1] = and(and.3206, not.3169, id=3207)
+  and.3241: bits[1] = and(and.3240, not.3166, id=3241)
+  and.3223: bits[1] = and(and.3222, not.3168, id=3223)
+  not.3213: bits[1] = not(and.3202, id=3213)
+  and.3226: bits[1] = and(bit_slice.2847, or.3212, id=3226)
+  and.3229: bits[1] = and(and.3228, bit_slice.2845, id=3229)
+  or.3216: bits[1] = or(not.3166, and.3207, id=3216)
+  and.3242: bits[1] = and(and.3241, not.3167, id=3242)
+  and.3224: bits[1] = and(and.3223, not.3169, id=3224)
+  bit_slice.2848: bits[1] = bit_slice(array_index.2844, start=3, width=1, id=2848)
+  or.3233: bits[1] = or(not.3211, and.3202, id=3233)
+  and.3225: bits[1] = and(or.3212, not.3213, id=3225)
+  and.3227: bits[1] = and(and.3226, not.3213, id=3227)
+  and.3230: bits[1] = and(and.3229, or.3216, id=3230)
+  not.3217: bits[1] = not(and.3178, id=3217)
+  and.3243: bits[1] = and(and.3242, not.3168, id=3243)
+  not.3234: bits[1] = not(and.3224, id=3234)
+  and.3246: bits[1] = and(bit_slice.2848, or.3233, id=3246)
+  or.3235: bits[1] = or(bit_slice.2847, and.3225, id=3235)
+  not.3236: bits[1] = not(and.3227, id=3236)
+  and.3231: bits[1] = and(and.3230, not.3217, id=3231)
+  and.3244: bits[1] = and(and.3243, not.3169, id=3244)
+  bit_slice.2849: bits[1] = bit_slice(array_index.2844, start=4, width=1, id=2849)
+  or.3250: bits[1] = or(not.3232, and.3224, id=3250)
+  and.3245: bits[1] = and(or.3233, not.3234, id=3245)
+  and.3247: bits[1] = and(and.3246, not.3234, id=3247)
+  and.3248: bits[1] = and(or.3235, not.3236, id=3248)
+  or.3237: bits[1] = or(and.3205, and.3231, id=3237)
+  not.3251: bits[1] = not(and.3244, id=3251)
+  and.3256: bits[1] = and(bit_slice.2849, or.3250, id=3256)
+  or.3252: bits[1] = or(bit_slice.2848, and.3245, id=3252)
+  not.3253: bits[1] = not(and.3247, id=3253)
+  and.3249: bits[1] = and(and.3248, or.3237, id=3249)
+  and.3255: bits[1] = and(or.3250, not.3251, id=3255)
+  and.3257: bits[1] = and(and.3256, not.3251, id=3257)
+  and.3258: bits[1] = and(or.3252, not.3253, id=3258)
+  or.3254: bits[1] = or(and.3227, and.3249, id=3254)
+  and.3266: bits[1] = and(or.3252, not.3253, id=3266)
+  or.3260: bits[1] = or(bit_slice.2849, and.3255, id=3260)
+  not.3261: bits[1] = not(and.3257, id=3261)
+  and.3259: bits[1] = and(and.3258, or.3254, id=3259)
+  or.3276: bits[1] = or(and.3266, and.3227, id=3276)
+  and.3267: bits[1] = and(or.3235, not.3236, id=3267)
+  and.3263: bits[1] = and(or.3260, not.3261, id=3263)
+  and.3264: bits[1] = and(or.3260, not.3261, id=3264)
+  or.3262: bits[1] = or(and.3247, and.3259, id=3262)
+  or.3277: bits[1] = or(or.3276, and.3249, id=3277)
+  not.3278: bits[1] = not(and.3259, id=3278)
+  or.3279: bits[1] = or(and.3267, and.3205, id=3279)
+  and.3269: bits[1] = and(bit_slice.2845, or.3216, id=3269)
+  and.3271: bits[1] = and(or.3216, not.3217, id=3271)
+  or.3273: bits[1] = or(and.3263, and.3247, id=3273)
+  and.3265: bits[1] = and(and.3264, or.3262, id=3265)
+  and.3292: bits[1] = and(or.3277, not.3278, id=3292)
+  or.3280: bits[1] = or(or.3279, and.3231, id=3280)
+  and.3268: bits[1] = and(or.3214, not.3215, id=3268)
+  and.3270: bits[1] = and(and.3269, not.3217, id=3270)
+  and.3298: bits[1] = and(or.3277, not.3278, id=3298)
+  or.3284: bits[1] = or(bit_slice.2845, and.3271, id=3284)
+  or.3274: bits[1] = or(or.3273, and.3259, id=3274)
+  not.3275: bits[1] = not(and.3265, id=3275)
+  and.3293: bits[1] = and(and.3292, or.3280, id=3293)
+  not.3281: bits[1] = not(and.3249, id=3281)
+  and.3295: bits[1] = and(or.3277, not.3278, id=3295)
+  or.3282: bits[1] = or(and.3268, and.3270, id=3282)
+  and.3299: bits[1] = and(and.3298, or.3284, id=3299)
+  not.3285: bits[1] = not(and.3270, id=3285)
+  and.3272: bits[1] = and(not.3168, not.3169, id=3272)
+  and.3288: bits[1] = and(or.3274, not.3275, id=3288)
+  and.3289: bits[1] = and(or.3277, not.3278, id=3289)
+  and.3294: bits[1] = and(and.3293, not.3281, id=3294)
+  and.3296: bits[1] = and(and.3295, or.3282, id=3296)
+  not.3283: bits[1] = not(and.3231, id=3283)
+  and.3300: bits[1] = and(and.3299, not.3285, id=3300)
+  or.3286: bits[1] = or(not.3167, and.3272, id=3286)
+  or.3303: bits[1] = or(and.3288, and.3289, id=3303)
+  and.3290: bits[1] = and(or.3280, not.3281, id=3290)
+  or.3311: bits[1] = or(and.3288, and.3294, id=3311)
+  and.3297: bits[1] = and(and.3296, not.3283, id=3297)
+  and.3301: bits[1] = and(and.3300, or.3286, id=3301)
+  not.3287: bits[1] = not(and.3207, id=3287)
+  or.3304: bits[1] = or(or.3303, and.3290, id=3304)
+  and.3291: bits[1] = and(or.3282, not.3283, id=3291)
+  or.3306: bits[1] = or(and.3288, and.3294, id=3306)
+  or.3312: bits[1] = or(or.3311, and.3297, id=3312)
+  and.3302: bits[1] = and(and.3301, not.3287, id=3302)
+  or.3305: bits[1] = or(or.3304, and.3291, id=3305)
+  or.3307: bits[1] = or(or.3306, and.3297, id=3307)
+  or.3313: bits[1] = or(or.3312, and.3302, id=3313)
+  not.3309: bits[1] = not(or.3305, id=3309)
+  or.3308: bits[1] = or(or.3307, and.3302, id=3308)
+  or.3318: bits[1] = or(not.3168, not.3169, id=3318)
+  not.3319: bits[1] = not(and.3272, id=3319)
+  or.3314: bits[1] = or(or.3313, not.3309, id=3314)
+  and.3322: bits[1] = and(or.3284, not.3285, id=3322)
+  not.3310: bits[1] = not(or.3308, id=3310)
+  and.3328: bits[1] = and(or.3286, not.3287, id=3328)
+  and.3333: bits[1] = and(or.3318, not.3319, id=3333)
+  and.3338: bits[1] = and(or.3138, not.3139, id=3338)
+  and.3320: bits[1] = and(or.3284, not.3285, id=3320)
+  not.3316: bits[1] = not(or.3314, id=3316)
+  and.3323: bits[1] = and(and.3322, not.3310, id=3323)
+  and.3315: bits[1] = and(not.3310, not.3309, id=3315)
+  and.3326: bits[1] = and(or.3286, not.3287, id=3326)
+  and.3329: bits[1] = and(and.3328, not.3310, id=3329)
+  and.3331: bits[1] = and(or.3318, not.3319, id=3331)
+  and.3334: bits[1] = and(and.3333, not.3310, id=3334)
+  and.3336: bits[1] = and(or.3138, not.3139, id=3336)
+  and.3339: bits[1] = and(and.3338, not.3310, id=3339)
+  and.3341: bits[1] = and(or.3282, not.3283, id=3341)
+  and.3321: bits[1] = and(and.3320, not.3316, id=3321)
+  and.3324: bits[1] = and(and.3323, not.3309, id=3324)
+  not.3317: bits[1] = not(and.3315, id=3317)
+  and.3327: bits[1] = and(and.3326, not.3316, id=3327)
+  and.3330: bits[1] = and(and.3329, not.3309, id=3330)
+  and.3332: bits[1] = and(and.3331, not.3316, id=3332)
+  and.3335: bits[1] = and(and.3334, not.3309, id=3335)
+  and.3337: bits[1] = and(and.3336, not.3316, id=3337)
+  and.3340: bits[1] = and(and.3339, not.3309, id=3340)
+  literal.2803: bits[4] = literal(value=0, id=2803)
+  literal.2823: bits[4] = literal(value=2, id=2823)
+  literal.2863: bits[4] = literal(value=6, id=2863)
+  literal.2883: bits[4] = literal(value=8, id=2883)
+  and.3342: bits[1] = and(and.3341, not.3310, id=3342)
+  or.3344: bits[1] = or(and.3321, and.3324, id=3344)
+  and.3325: bits[1] = and(not.3316, not.3317, id=3325)
+  or.3346: bits[1] = or(and.3327, and.3330, id=3346)
+  or.3348: bits[1] = or(and.3332, and.3335, id=3348)
+  or.3350: bits[1] = or(and.3337, and.3340, id=3350)
+  array_index.2804: bits[8] = array_index(window, indices=[literal.2803], id=2804)
+  array_index.2824: bits[8] = array_index(window, indices=[literal.2823], id=2824)
+  array_index.2864: bits[8] = array_index(window, indices=[literal.2863], id=2864)
+  array_index.2884: bits[8] = array_index(window, indices=[literal.2883], id=2884)
+  literal.2801: bits[1] = literal(value=0, id=2801)
+  and.3343: bits[1] = and(and.3342, not.3309, id=3343)
+  or.3345: bits[1] = or(or.3344, and.3325, id=3345)
+  or.3347: bits[1] = or(or.3346, and.3325, id=3347)
+  or.3349: bits[1] = or(or.3348, and.3325, id=3349)
+  or.3351: bits[1] = or(or.3350, and.3325, id=3351)
+  literal.2800: bits[1] = literal(value=1, id=2800)
+  bit_slice.2805: bits[1] = bit_slice(array_index.2804, start=0, width=1, id=2805)
+  bit_slice.2806: bits[1] = bit_slice(array_index.2804, start=1, width=1, id=2806)
+  bit_slice.2807: bits[1] = bit_slice(array_index.2804, start=2, width=1, id=2807)
+  bit_slice.2808: bits[1] = bit_slice(array_index.2804, start=3, width=1, id=2808)
+  bit_slice.2809: bits[1] = bit_slice(array_index.2804, start=4, width=1, id=2809)
+  bit_slice.2810: bits[1] = bit_slice(array_index.2804, start=5, width=1, id=2810)
+  bit_slice.2811: bits[1] = bit_slice(array_index.2804, start=6, width=1, id=2811)
+  bit_slice.2812: bits[1] = bit_slice(array_index.2804, start=7, width=1, id=2812)
+  bit_slice.2815: bits[1] = bit_slice(array_index.2814, start=0, width=1, id=2815)
+  bit_slice.2825: bits[1] = bit_slice(array_index.2824, start=0, width=1, id=2825)
+  bit_slice.2826: bits[1] = bit_slice(array_index.2824, start=1, width=1, id=2826)
+  bit_slice.2827: bits[1] = bit_slice(array_index.2824, start=2, width=1, id=2827)
+  bit_slice.2828: bits[1] = bit_slice(array_index.2824, start=3, width=1, id=2828)
+  bit_slice.2829: bits[1] = bit_slice(array_index.2824, start=4, width=1, id=2829)
+  bit_slice.2830: bits[1] = bit_slice(array_index.2824, start=5, width=1, id=2830)
+  bit_slice.2831: bits[1] = bit_slice(array_index.2824, start=6, width=1, id=2831)
+  bit_slice.2832: bits[1] = bit_slice(array_index.2824, start=7, width=1, id=2832)
+  bit_slice.2835: bits[1] = bit_slice(array_index.2834, start=0, width=1, id=2835)
+  bit_slice.2850: bits[1] = bit_slice(array_index.2844, start=5, width=1, id=2850)
+  bit_slice.2851: bits[1] = bit_slice(array_index.2844, start=6, width=1, id=2851)
+  bit_slice.2852: bits[1] = bit_slice(array_index.2844, start=7, width=1, id=2852)
+  bit_slice.2855: bits[1] = bit_slice(array_index.2854, start=0, width=1, id=2855)
+  bit_slice.2865: bits[1] = bit_slice(array_index.2864, start=0, width=1, id=2865)
+  bit_slice.2866: bits[1] = bit_slice(array_index.2864, start=1, width=1, id=2866)
+  bit_slice.2867: bits[1] = bit_slice(array_index.2864, start=2, width=1, id=2867)
+  bit_slice.2868: bits[1] = bit_slice(array_index.2864, start=3, width=1, id=2868)
+  bit_slice.2869: bits[1] = bit_slice(array_index.2864, start=4, width=1, id=2869)
+  bit_slice.2870: bits[1] = bit_slice(array_index.2864, start=5, width=1, id=2870)
+  bit_slice.2871: bits[1] = bit_slice(array_index.2864, start=6, width=1, id=2871)
+  bit_slice.2872: bits[1] = bit_slice(array_index.2864, start=7, width=1, id=2872)
+  bit_slice.2875: bits[1] = bit_slice(array_index.2874, start=0, width=1, id=2875)
+  bit_slice.2885: bits[1] = bit_slice(array_index.2884, start=0, width=1, id=2885)
+  bit_slice.2886: bits[1] = bit_slice(array_index.2884, start=1, width=1, id=2886)
+  bit_slice.2887: bits[1] = bit_slice(array_index.2884, start=2, width=1, id=2887)
+  bit_slice.2888: bits[1] = bit_slice(array_index.2884, start=3, width=1, id=2888)
+  bit_slice.2889: bits[1] = bit_slice(array_index.2884, start=4, width=1, id=2889)
+  bit_slice.2890: bits[1] = bit_slice(array_index.2884, start=5, width=1, id=2890)
+  bit_slice.2891: bits[1] = bit_slice(array_index.2884, start=6, width=1, id=2891)
+  bit_slice.2892: bits[1] = bit_slice(array_index.2884, start=7, width=1, id=2892)
+  ret concat.3352: bits[8] = concat(literal.2801, literal.2801, literal.2801, and.3343, or.3345, or.3347, or.3349, or.3351, id=3352)
+}

cccl_upstream/cudax/test/stf/fhe/parse_arctyrex.cu

@@ -0,0 +1,434 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief This shows how we can asynchronously compose a sequence of operations
+ *        described with an IR
+ */
+
+#include <cuda/experimental/stf.cuh>
+
+#include <fstream>
+#include <iostream>
+#include <map>
+
+using namespace cuda::experimental::stf;
+
+using logical_slice = logical_data<slice<double>>;
+
+static __global__ void cuda_sleep_kernel(long long int clock_cnt)
+{
+  long long int start_clock  = clock64();
+  long long int clock_offset = 0;
+  while (clock_offset < clock_cnt)
+  {
+    clock_offset = clock64() - start_clock;
+  }
+}
+
+void cuda_sleep(double ms, cudaStream_t stream)
+{
+  int device;
+  cudaGetDevice(&device);
+
+  // cudaDevAttrClockRate: Peak clock frequency in kilohertz;
+  int clock_rate;
+  cudaDeviceGetAttribute(&clock_rate, cudaDevAttrClockRate, device);
+
+  long long int clock_cnt = (long long int) (ms * clock_rate);
+  cuda_sleep_kernel<<<1, 1, 0, stream>>>(clock_cnt);
+}
+
+const double sleep_time = 1.0;
+
+// z = LITERAL(length, value)
+template <typename Ctx>
+logical_slice LITERAL(Ctx& ctx, size_t n, int, std::string out_symbol = "undefined")
+{
+  auto z = ctx.logical_data(shape_of<slice<double>>(n));
+  z.set_symbol(out_symbol);
+
+  ctx.task(z.write()).set_symbol("LITERAL")->*[](cudaStream_t stream, auto /*unused*/) {
+    cuda_sleep(sleep_time, stream);
+  };
+
+  return z;
+}
+
+// z = OR(x,y)
+template <typename Ctx>
+logical_slice OR(Ctx& ctx, logical_slice x, logical_slice y, std::string out_symbol = "undefined")
+{
+  assert(x.shape().size() == y.shape().size());
+
+  auto z = ctx.logical_data(x.shape());
+  z.set_symbol(out_symbol);
+
+  ctx.task(x.read(), y.read(), z.write()).set_symbol("OR")->*
+    [](cudaStream_t stream, auto /*unused*/, auto /*unused*/, auto /*unused*/) {
+      cuda_sleep(sleep_time, stream);
+    };
+
+  return z;
+}
+
+// z = AND(x,y)
+template <typename Ctx>
+logical_slice AND(Ctx& ctx, logical_slice x, logical_slice y, std::string out_symbol = "undefined")
+{
+  assert(x.shape().size() == y.shape().size());
+
+  auto z = ctx.logical_data(x.shape());
+  z.set_symbol(out_symbol);
+
+  ctx.task(x.read(), y.read(), z.write()).set_symbol("AND")->*
+    [](cudaStream_t stream, auto /*unused*/, auto /*unused*/, auto /*unused*/) {
+      cuda_sleep(sleep_time, stream);
+    };
+
+  return z;
+}
+
+template <typename Ctx>
+logical_slice
+ARRAY_INDEX(Ctx& ctx, logical_slice x, logical_slice /*unused*/, size_t sz, std::string out_symbol = "undefined")
+{
+  auto z = ctx.logical_data(shape_of<slice<double>>(sz));
+  z.set_symbol(out_symbol);
+
+  ctx.task(x.read(), z.write()).set_symbol("ARRAY INDEX")->*[](cudaStream_t stream, auto /*unused*/, auto /*unused*/) {
+    cuda_sleep(sleep_time, stream);
+  };
+
+  return z;
+}
+
+// z = BIT_SLICE(x, position, size)
+template <typename Ctx>
+logical_slice BIT_SLICE(Ctx& ctx, logical_slice x, size_t /*unused*/, size_t sz, std::string out_symbol = "undefined")
+{
+  auto z = ctx.logical_data(shape_of<slice<double>>(sz));
+  z.set_symbol(out_symbol);
+
+  ctx.task(x.read(), z.write()).set_symbol("BIT SLICE")->*[](cudaStream_t stream, auto /*unused*/, auto /*unused*/) {
+    cuda_sleep(sleep_time, stream);
+  };
+
+  return z;
+}
+
+// y = NOT(x)
+template <typename Ctx>
+logical_slice NOT(Ctx& ctx, logical_slice x, std::string out_symbol = "undefined")
+{
+  auto y = ctx.logical_data(x.shape());
+  y.set_symbol(out_symbol);
+
+  ctx.task(x.read(), y.write()).set_symbol("NOT")->*[](cudaStream_t stream, auto /*unused*/, auto /*unused*/) {
+    cuda_sleep(sleep_time, stream);
+  };
+
+  return y;
+}
+
+// y = CONCAT(sz, vector<> inputs)
+template <typename Ctx>
+logical_slice CONCAT(Ctx& ctx, size_t sz, std::vector<logical_slice> inputs, std::string out_symbol = "undefined")
+{
+  auto y = ctx.logical_data(shape_of<slice<double>>(sz));
+  y.set_symbol(out_symbol);
+
+  auto t = ctx.task();
+  t.add_deps(y.write());
+  t.set_symbol("CONCAT");
+  for (auto& input : inputs)
+  {
+    t.add_deps(input.read());
+  }
+
+  t->*[](cudaStream_t stream) {
+    cuda_sleep(sleep_time, stream);
+  };
+
+  return y;
+}
+
+template <typename Ctx>
+void run(const char* inputfile)
+{
+  // Find the handle from its symbol
+  std::map<std::string, logical_slice> logical_slices;
+  std::string output_data_symbol;
+  std::ifstream read(inputfile);
+  Ctx ctx;
+
+  // Indicates if we are parsing the body of the circuit
+  bool in_body = false;
+
+  for (std::string line; std::getline(read, line);)
+  {
+    std::stringstream ss(line);
+
+    //        std::cout << "LINE : " << line << '\n';
+    if (!in_body)
+    {
+      std::string token;
+      ss >> token;
+      //            std::cout << "TOKEN : " << token << '\n';
+
+      if (token == "fn")
+      {
+        // We are parsing the declaration of the function, this starts the body
+        // std::cout << "GOT DECLARATION " << line << '\n';
+        in_body = true;
+
+        // Look for parameters
+        size_t begin_params, end_params;
+        begin_params       = line.find('(');
+        end_params         = line.find(')');
+        std::string params = line.substr(begin_params + 1, end_params - begin_params - 1);
+
+        // std::cout << "PARAMS = " << params << '\n';
+
+        // Parse parameters which are separated by a comma, format = "symbol: type"
+        while (true)
+        {
+          // Find symbol
+          size_t pos;
+          pos                = params.find(":");
+          std::string symbol = params.substr(0, pos);
+          // std::cout << symbol << '\n';
+
+          // We create a dummy allocation so that the data handles refers to actually allocated host memory
+          double* dummy     = new double[1];
+          auto param_handle = ctx.logical_data(make_slice(dummy, 1));
+
+          param_handle.set_symbol(symbol);
+          logical_slices[symbol] = param_handle;
+
+          pos = params.find(", ");
+          if (pos == std::string::npos)
+          {
+            break;
+          }
+
+          params.erase(0, pos + 2);
+        }
+      }
+    }
+    else
+    {
+      std::string token;
+      ss >> token;
+      // std::cout << "TOKEN : " << token << '\n';
+
+      if (token == "}")
+      {
+        // This closes the body
+        in_body = false;
+        continue;
+      }
+
+      // We expect lines of the format : "  symbol: type = gate_name(..., id=VALUE)"
+      size_t end_symbol = line.find(":");
+
+      // We look for the first "= " to find the gate name
+      size_t gate_symbol_pos = line.find("= ");
+      std::string gate       = line.substr(gate_symbol_pos + 2);
+
+      size_t gate_name_end           = gate.find("(");
+      std::string gate_symbol        = gate.substr(0, gate_name_end);
+      std::string gate_args          = gate.substr(gate_name_end + 1, gate.size() - gate_name_end - 2);
+      std::string gate_outvar_symbol = line.substr(2, end_symbol - 2);
+
+      // Possibly remove the "ret" out of the gate_outvar_symbol
+      size_t ret_pos = gate_outvar_symbol.find("ret ");
+      if (ret_pos != std::string::npos)
+      {
+        // This is our result !
+        gate_outvar_symbol.erase(4);
+
+        output_data_symbol = gate_outvar_symbol;
+      }
+
+      // std::cout << "GATE OUT SYMBOL " << gate_outvar_symbol << '\n';
+      // std::cout << "GATE DESCRIPTION : " << gate << '\n';
+      // std::cout << "GATE SYMBOL " << gate_symbol << '\n';
+      // std::cout << "GATE ARGS " << gate_args << '\n';
+
+      // We now dispatch between the different gates
+      if (gate_symbol == "literal")
+      {
+        //  literal.916: bits[1] = literal(value=1, id=916)
+        int value = 42; // TODO parse
+        size_t sz = 1;
+
+        logical_slices[gate_outvar_symbol] = LITERAL(ctx, sz, value, gate_outvar_symbol);
+
+        continue;
+      }
+
+      if (gate_symbol == "or")
+      {
+        //  or.1268: bits[1] = or(or.1267, and.1236, id=1268)
+        size_t pos;
+        pos                     = gate_args.find(", ");
+        std::string symbol_left = gate_args.substr(0, pos);
+        gate_args.erase(0, pos + 2);
+
+        pos                      = gate_args.find(", ");
+        std::string symbol_right = gate_args.substr(0, pos);
+        gate_args.erase(0, pos + 2);
+
+        // std::cout << "OR GATE on symbols" << symbol_left << " AND " << symbol_right << '\n';
+
+        auto data_left                     = logical_slices[symbol_left];
+        auto data_right                    = logical_slices[symbol_right];
+        logical_slices[gate_outvar_symbol] = OR(ctx, data_left, data_right, gate_outvar_symbol);
+
+        continue;
+      }
+
+      if (gate_symbol == "and")
+      {
+        size_t pos;
+        pos                     = gate_args.find(", ");
+        std::string symbol_left = gate_args.substr(0, pos);
+        gate_args.erase(0, pos + 2);
+
+        pos                      = gate_args.find(", ");
+        std::string symbol_right = gate_args.substr(0, pos);
+        gate_args.erase(0, pos + 2);
+
+        auto data_left                     = logical_slices[symbol_left];
+        auto data_right                    = logical_slices[symbol_right];
+        logical_slices[gate_outvar_symbol] = AND(ctx, data_left, data_right, gate_outvar_symbol);
+
+        continue;
+      }
+
+      if (gate_symbol == "bit_slice")
+      {
+        // bit_slice.936: bits[1] = bit_slice(y, start=15, width=1, id=936)
+        size_t pos;
+        pos                   = gate_args.find(", ");
+        std::string symbol_in = gate_args.substr(0, pos);
+        gate_args.erase(0, pos + 2);
+
+        // hardcoded ...
+        size_t sz = 1;
+
+        auto data_in                       = logical_slices[symbol_in];
+        logical_slices[gate_outvar_symbol] = BIT_SLICE(ctx, data_in, 42, sz, gate_outvar_symbol);
+
+        // std::cout << "PRODUCED DATA FOR " << gate_outvar_symbol << '\n';
+
+        continue;
+      }
+
+      if (gate_symbol == "array_index")
+      {
+        // array_index.2804: bits[8] = array_index(window, indices=[literal.2803], id=2804)
+        size_t pos;
+        pos                   = gate_args.find(", ");
+        std::string symbol_in = gate_args.substr(0, pos);
+        gate_args.erase(0, pos + 2);
+
+        pos                        = gate_args.find(", ");
+        std::string symbol_indices = gate_args.substr(0, pos);
+        gate_args.erase(0, pos + 2);
+
+        size_t pos_beg          = symbol_indices.find("[");
+        size_t pos_end          = symbol_indices.find("]");
+        std::string symbol_in_2 = symbol_indices.substr(pos_beg + 1, pos_end - pos_beg - 1);
+        // std::cout << "ARRAY INDEX ... INDEX = " << symbol_in_2 << '\n';
+
+        // hardcoded ...
+        size_t sz = 1;
+
+        auto data_in                       = logical_slices[symbol_in];
+        auto data_in_2                     = logical_slices[symbol_in_2];
+        logical_slices[gate_outvar_symbol] = ARRAY_INDEX(ctx, data_in, data_in_2, sz, gate_outvar_symbol);
+
+        // std::cout << "PRODUCED DATA FOR " << gate_outvar_symbol << '\n';
+        continue;
+      }
+
+      if (gate_symbol == "not")
+      {
+        // not.953: bits[1] = not(bit_slice.936, id=953)
+        size_t pos;
+        pos                   = gate_args.find(", ");
+        std::string symbol_in = gate_args.substr(0, pos);
+        gate_args.erase(0, pos + 2);
+
+        auto data_in                       = logical_slices[symbol_in];
+        logical_slices[gate_outvar_symbol] = NOT(ctx, data_in, gate_outvar_symbol);
+
+        continue;
+      }
+
+      if (gate_symbol == "concat")
+      {
+        //   ret concat.1269: bits[16] = concat(or.1238, or.1240, or.1242, or.1244, or.1246, or.1248, or.1250,
+        //   or.1252, or.1254, or.1256, or.1258, or.1260, or.1262, or.1264, or.1266, or.1268, id=1269)
+        // Remove the end ", id =.."
+        size_t id_pos = gate_args.find(", id=");
+        gate_args     = gate_args.substr(0, id_pos);
+
+        std::vector<logical_slice> inputs;
+        size_t pos;
+
+        while (true)
+        {
+          pos = gate_args.find(", ");
+          if (pos == std::string::npos)
+          {
+            break;
+          }
+
+          std::string symbol = gate_args.substr(0, pos);
+          inputs.push_back(logical_slices[symbol]);
+          gate_args.erase(0, pos + 2);
+
+          // std::cout << "CONCAT ARG = " << symbol << '\n';
+        }
+
+        size_t sz                          = 1;
+        logical_slices[gate_outvar_symbol] = CONCAT(ctx, sz, inputs, gate_outvar_symbol);
+        continue;
+      }
+
+      std::cout << "UNRECOGNIZED GATE !" << '\n';
+      abort();
+    }
+  }
+
+  auto output_data = logical_slices[output_data_symbol];
+
+  ctx.finalize();
+}
+
+int main(int argc, char** argv)
+{
+  /* Until we find a simple and "safe" way to pass a file to the test, we
+   * consider this is not an error ...
+   * One possible approach would be to convert a default .ir file to a large
+   * static data processed in the test suite ?
+   */
+  if (argc < 2)
+  {
+    fprintf(stderr, "This test needs an input file, skipping.\n");
+    return 0;
+  }
+  run<stream_ctx>(argv[1]);
+  run<graph_ctx>(argv[1]);
+}

cccl_upstream/cudax/test/stf/freeze/constant_logical_data.cu

@@ -0,0 +1,72 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ *
+ * @brief Helper to build constant data based on frozen logical data
+ *
+ */
+
+#include <cuda/experimental/__stf/utility/constant_logical_data.cuh>
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+int main()
+{
+  stream_ctx ctx;
+
+  const int N = 16;
+
+  /* Create a constant value */
+  auto ld_cst = ctx.logical_data(shape_of<slice<int>>(N));
+  ctx.parallel_for(ld_cst.shape(), ld_cst.write())->*[] __device__(size_t i, slice<int> res) {
+    res(i) = 18 * i - 9;
+  };
+  auto cst = constant_logical_data(ctx, mv(ld_cst));
+
+  int X[N];
+
+  for (int i = 0; i < N; i++)
+  {
+    X[i] = 5 * i - 3;
+  }
+
+  auto lX = ctx.logical_data(X).set_symbol("X");
+
+  for (size_t iter = 0; iter < 4; iter++)
+  {
+    auto cst_slice = cst.get();
+    ctx.parallel_for(lX.shape(), lX.rw()).set_symbol("X+=cst")->*[cst_slice] __device__(size_t i, auto x) {
+      x(i) += cst_slice(i);
+    };
+
+    auto cst2 = run_once()->*[&]() {
+      auto ld = ctx.logical_data(shape_of<slice<int>>(N));
+      ctx.parallel_for(ld.shape(), ld.write())->*[] __device__(size_t i, slice<int> res) {
+        res(i) = 4 * i - 2;
+      };
+      return constant_logical_data(ctx, mv(ld));
+    };
+
+    auto cst2_slice = cst2.get();
+    ctx.parallel_for(lX.shape(), lX.rw()).set_symbol("X+=cst2")->*[cst2_slice] __device__(size_t i, auto x) {
+      x(i) += cst2_slice(i);
+    };
+  }
+
+  ctx.finalize();
+
+  for (int i = 0; i < N; i++)
+  {
+    EXPECT(X[i] == (5 * i - 3) + 4 * (18 * i - 9) + 4 * (4 * i - 2));
+  }
+}

cccl_upstream/cudax/test/stf/freeze/freeze.cu

@@ -0,0 +1,82 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ *
+ * @brief Freeze data in read-only fashion
+ *
+ */
+
+#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+int X0(int i)
+{
+  return 17 * i + 45;
+}
+
+__global__ void print(slice<int> s)
+{
+  int tid      = blockIdx.x * blockDim.x + threadIdx.x;
+  int nthreads = gridDim.x * blockDim.x;
+
+  for (int i = tid; i < s.size(); i += nthreads)
+  {
+    printf("%d %d\n", i, s(i));
+  }
+}
+
+int main()
+{
+  stream_ctx ctx;
+
+  cudaStream_t stream = ctx.pick_stream();
+
+  const int N = 16;
+  int X[N];
+
+  for (int i = 0; i < N; i++)
+  {
+    X[i] = X0(i);
+  }
+
+  auto lX = ctx.logical_data(X).set_symbol("X");
+  auto lY = ctx.logical_data(lX.shape()).set_symbol("Y");
+
+  ctx.parallel_for(lX.shape(), lX.rw()).set_symbol("X=2X")->*[] __device__(size_t i, auto x) {
+    x(i) *= 2;
+  };
+
+  auto fx = ctx.freeze(lX);
+
+  auto dX = fx.get(data_place::current_device(), stream);
+
+  print<<<8, 4, 0, stream>>>(dX);
+
+  ctx.parallel_for(lX.shape(), lX.read(), lY.write()).set_symbol("Y=X")->*[] __device__(size_t i, auto x, auto y) {
+    y(i) = x(i);
+  };
+
+  fx.unfreeze(stream);
+
+  ctx.parallel_for(lX.shape(), lX.rw()).set_symbol("X+=1")->*[] __device__(size_t i, auto x) {
+    x(i) += 1;
+  };
+
+  ctx.finalize();
+
+  for (int i = 0; i < N; i++)
+  {
+    EXPECT(X[i] == 2 * X0(i) + 1);
+  }
+}

cccl_upstream/cudax/test/stf/freeze/freeze_rw.cu

@@ -0,0 +1,86 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ *
+ * @brief Freeze data in read-only fashion
+ *
+ */
+
+#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+int X0(int i)
+{
+  return 17 * i + 45;
+}
+
+__global__ void mult(slice<int> s, int val)
+{
+  int tid      = blockIdx.x * blockDim.x + threadIdx.x;
+  int nthreads = gridDim.x * blockDim.x;
+
+  for (int i = tid; i < s.size(); i += nthreads)
+  {
+    s(i) *= val;
+  }
+}
+
+int main()
+{
+  stream_ctx ctx;
+
+  cudaStream_t stream = ctx.pick_stream();
+
+  const int N = 16;
+  int X[N];
+
+  for (int i = 0; i < N; i++)
+  {
+    X[i] = X0(i);
+  }
+
+  auto lX = ctx.logical_data(X).set_symbol("X");
+  auto lY = ctx.logical_data(lX.shape()).set_symbol("Y");
+
+  for (int k = 0; k < 4; k++)
+  {
+    auto fx = ctx.freeze(lX, access_mode::rw, data_place::current_device());
+
+    _CCCL_ASSERT(fx.get_access_mode() == access_mode::rw, "invalid access mode");
+
+    auto dX = fx.get(data_place::current_device(), stream);
+    mult<<<8, 4, 0, stream>>>(dX, 4);
+    fx.unfreeze(stream);
+
+    ctx.parallel_for(lX.shape(), lX.read(), lY.write()).set_symbol("Y=X")->*[] __device__(size_t i, auto x, auto y) {
+      y(i) = x(i);
+    };
+
+    ctx.parallel_for(lX.shape(), lY.rw()).set_symbol("Y+=1")->*[] __device__(size_t i, auto y) {
+      y(i) += 1;
+    };
+
+    //        ctx.host_launch(lX.read(), lY.read())->*[](auto x, auto y) {
+    //            for (int i = 0; i < x.size(); i++) {
+    //                EXPECT(x(i) == 2*X0(i) + 4);
+    //            }
+    //
+    //            for (int i = 0; i < y.size(); i++) {
+    //                EXPECT(y(i) == 2*X0(i) + 4);
+    //            }
+    //        };
+  }
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/freeze/freeze_untyped_rw.cu

@@ -0,0 +1,74 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+//! \file
+//!
+//! \brief Freeze data and store it as a frozen_logical_data_untyped object
+
+#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+int X0(int i)
+{
+  return 17 * i + 45;
+}
+
+__global__ void mult(slice<int> s, int val)
+{
+  int tid      = blockIdx.x * blockDim.x + threadIdx.x;
+  int nthreads = gridDim.x * blockDim.x;
+
+  for (int i = tid; i < s.size(); i += nthreads)
+  {
+    s(i) *= val;
+  }
+}
+
+int main()
+{
+  stream_ctx ctx;
+
+  cudaStream_t stream = ctx.pick_stream();
+
+  const int N = 16;
+  int X[N];
+
+  for (int i = 0; i < N; i++)
+  {
+    X[i] = X0(i);
+  }
+
+  auto lX = ctx.logical_data(X).set_symbol("X");
+  auto lY = ctx.logical_data(lX.shape()).set_symbol("Y");
+
+  for (int k = 0; k < 4; k++)
+  {
+    logical_data_untyped lX_untyped = lX;
+    auto fx                         = ctx.freeze(lX_untyped, access_mode::rw, data_place::current_device());
+
+    _CCCL_ASSERT(fx.get_access_mode() == access_mode::rw, "invalid access mode");
+
+    auto dX = fx.template get<slice<int>>(data_place::current_device(), stream);
+    mult<<<8, 4, 0, stream>>>(dX, 4);
+    fx.unfreeze(stream);
+
+    ctx.parallel_for(lX.shape(), lX.read(), lY.write()).set_symbol("Y=X")->*[] __device__(size_t i, auto x, auto y) {
+      y(i) = x(i);
+    };
+
+    ctx.parallel_for(lX.shape(), lY.rw()).set_symbol("Y+=1")->*[] __device__(size_t i, auto y) {
+      y(i) += 1;
+    };
+  }
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/freeze/freeze_write_back.cu

@@ -0,0 +1,65 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+//! \file
+//!
+//! \brief Ensure write-back is working on logical data alias made by freezing another one
+
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+int main()
+{
+  context ctx;
+
+  int array[1024];
+  for (size_t i = 0; i < 1024; i++)
+  {
+    array[i] = 2 - i * i;
+  }
+
+  auto lA = ctx.logical_data(array).set_symbol("A");
+
+  auto stream = ctx.pick_stream();
+
+  graph_ctx gctx(stream);
+
+  // Create an alias for lA in the graph by freezing it and creating a new
+  // local logical data in the graph.
+  auto fa   = ctx.freeze(lA, access_mode::rw, data_place::current_device());
+  auto inst = fa.get(data_place::current_device(), stream);
+  auto glA  = gctx.logical_data(inst, data_place::current_device());
+
+  gctx.parallel_for(glA.shape(), glA.rw())->*[] __device__(size_t i, auto a) {
+    a(i) += 4 * i;
+  };
+
+  // force to move to a different place, and probably to allocate another copy
+  // on the host. This tests if the write-back mechanism works from the host to
+  // the device when destroying the alias logical data glA.
+  gctx.host_launch(glA.rw())->*[](auto a) {
+    for (size_t i = 0; i < 1024; i++)
+    {
+      a(i) *= 2;
+    }
+  };
+
+  gctx.finalize();
+
+  fa.unfreeze(stream);
+
+  ctx.finalize();
+
+  for (size_t i = 0; i < 1024; i++)
+  {
+    EXPECT(array[i] == 2 * (2 - i * i + 4 * i));
+  }
+}

cccl_upstream/cudax/test/stf/freeze/task_fence.cu

@@ -0,0 +1,98 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ *
+ * @brief Freeze data in read-only fashion
+ *
+ */
+
+#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+int X0(int i)
+{
+  return 17 * i + 45;
+}
+
+__global__ void print(slice<int> s)
+{
+  int tid      = blockIdx.x * blockDim.x + threadIdx.x;
+  int nthreads = gridDim.x * blockDim.x;
+
+  for (int i = tid; i < s.size(); i += nthreads)
+  {
+    printf("%d %d\n", i, s(i));
+  }
+}
+
+int main()
+{
+  stream_ctx ctx;
+
+  const int N = 16;
+  int X[N];
+
+  for (int i = 0; i < N; i++)
+  {
+    X[i] = X0(i);
+  }
+
+  auto lX = ctx.logical_data(X).set_symbol("X");
+  auto lY = ctx.logical_data(lX.shape()).set_symbol("Y");
+
+  ctx.parallel_for(lX.shape(), lX.rw()).set_symbol("X=2X")->*[] __device__(size_t i, auto x) {
+    x(i) *= 2;
+  };
+
+  // test 1 : implicit sync of gets
+  {
+    auto fx      = ctx.freeze(lX);
+    auto [dX, _] = fx.get(data_place::current_device());
+
+    // the stream returned by fence should depend on the get operation
+    auto stream2 = ctx.fence();
+    print<<<8, 4, 0, stream2>>>(dX);
+
+    ctx.parallel_for(lX.shape(), lX.read(), lY.write()).set_symbol("Y=X")->*[] __device__(size_t i, auto x, auto y) {
+      y(i) = x(i);
+    };
+
+    fx.unfreeze(stream2);
+  }
+
+  // test 2 : unfreeze with no events due to user sync
+  {
+    auto fx      = ctx.freeze(lX);
+    auto [dX, _] = fx.get(data_place::current_device());
+
+    // the stream returned by fence should depend on the get operation
+    auto stream2 = ctx.fence();
+    print<<<8, 4, 0, stream2>>>(dX);
+
+    // We synchronize so there is nothing to depend on anymore
+    cudaStreamSynchronize(stream2);
+    fx.unfreeze(event_list());
+  }
+
+  ctx.parallel_for(lX.shape(), lX.rw()).set_symbol("X+=1")->*[] __device__(size_t i, auto x) {
+    x(i) += 1;
+  };
+
+  ctx.finalize();
+
+  for (int i = 0; i < N; i++)
+  {
+    EXPECT(X[i] == 2 * X0(i) + 1);
+  }
+}

cccl_upstream/cudax/test/stf/freeze/token.cu

@@ -0,0 +1,37 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ *
+ * @brief Freeze token
+ *
+ */
+
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+int main()
+{
+  context ctx;
+
+  auto ltoken = ctx.token();
+
+  auto ftoken = ctx.freeze(ltoken);
+
+  cudaStream_t stream = ctx.pick_stream();
+  // This makes any future operations in this CUDA stream depend on the
+  // availability of the token
+  [[maybe_unused]] auto dtoken = ftoken.get(data_place::current_device(), stream);
+  ftoken.unfreeze(stream);
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/gnu/06-pdgemm.cpp

@@ -0,0 +1,366 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief An example that implements a tiled matrix product over multiple devices using CUBLAS
+ *
+ * This also illustrates how the same code base can be used both with a
+ * stream_ctx and a graph_ctx backend.
+ */
+
+#include <cuda/experimental/__stf/utility/nvtx.cuh>
+#include <cuda/experimental/stf.cuh>
+
+#define TILED
+
+using namespace cuda::experimental::stf;
+
+static std::unordered_map<exec_place, cublasHandle_t, hash<exec_place>> cublas_handles;
+
+/* Get a CUBLAS handle valid on the current execution place, or initialize it lazily */
+cublasHandle_t get_cublas_handle(const exec_place& ep = exec_place::current_device())
+{
+  auto& result = cublas_handles[ep];
+  if (result == cublasHandle_t())
+  { // not found, default value inserted
+    // Lazy initialization, and save the handle for future use
+    cuda_safe_call(cublasCreate(&result));
+  }
+  return result;
+}
+
+template <typename T>
+class matrix
+{
+public:
+  template <typename Ctx>
+  matrix(
+    Ctx& ctx, size_t NROWS, size_t NCOLS, size_t BLOCKSIZE_ROWS, size_t BLOCKSIZE_COLS, const char* _symbol = "matrix")
+  {
+    symbol = _symbol;
+
+    m  = NROWS;
+    mb = BLOCKSIZE_ROWS;
+
+    n  = NCOLS;
+    nb = BLOCKSIZE_COLS;
+
+    assert(m % mb == 0);
+    assert(n % nb == 0);
+
+    size_t s = ((size_t) m) * ((size_t) n) * sizeof(T);
+    // cuda_safe_call(cudaMallocHost(&h_array, m*n*sizeof(T)));
+    // fprintf(stderr, "Allocating %ld x %ld x %ld = %ld bytes (%f GB) on host for %s\n", m, n, sizeof(T), s,
+    //        s / (1024.0 * 1024.0 * 1024.0), _symbol);
+    h_array = (T*) malloc(s);
+    assert(h_array);
+    cuda_safe_call(cudaHostRegister(h_array, s, cudaHostRegisterPortable));
+
+    // Compute the number of blocks
+    mt = m / mb;
+    nt = n / nb;
+
+    handles.resize(mt * nt);
+
+    for (size_t colb = 0; colb < nt; colb++)
+    {
+      for (size_t rowb = 0; rowb < mt; rowb++)
+      {
+        T* addr_h = get_block_h(rowb, colb);
+
+#ifdef TILED
+        // tiles are stored contiguously
+        const size_t ld = mb;
+#else
+        const size_t ld = m;
+#endif
+
+        std::ignore = ld; // avoid warning #177-D: variable "ld" was declared but never referenced
+        auto s      = make_slice(addr_h, std::tuple{mb, nb}, ld);
+        auto tile   = ctx.logical_data(s);
+
+        tile.set_symbol(std::string(symbol) + "_" + std::to_string(rowb) + "_" + std::to_string(colb));
+
+        handles[rowb + colb * mt] = std::move(tile);
+      }
+    }
+
+    cuda_safe_call(cudaGetDeviceCount(&ndevs));
+    for (int a = 1; a * a <= ndevs; a++)
+    {
+      if (ndevs % a == 0)
+      {
+        grid_p = a;
+        grid_q = ndevs / a;
+      }
+    }
+
+    assert(grid_p * grid_q == ndevs);
+
+    // std::cout << "FOUND " << ndevs << " DEVICES "
+    //           << "p=" << grid_p << " q=" << grid_q << '\n';
+  }
+
+  int get_preferred_devid(int row, int col)
+  {
+    return (row % grid_p) + (col % grid_q) * grid_p;
+  }
+
+  logical_data<slice<T, 2>>& get_handle(int row, int col)
+  {
+    return handles[row + col * mt];
+  }
+
+  size_t get_index(size_t row, size_t col)
+  {
+#ifdef TILED
+    // Find which tile contains this element
+    int tile_row = row / mb;
+    int tile_col = col / nb;
+
+    size_t tile_size = mb * nb;
+
+    // Look for the index of the beginning of the tile
+    size_t tile_start = (tile_row + mt * tile_col) * tile_size;
+
+    // Offset within the tile
+    size_t offset = (row % mb) + (col % nb) * mb;
+
+    return tile_start + offset;
+#else
+    return row + col * m;
+#endif
+  }
+
+  T* get_block_h(int brow, int bcol)
+  {
+    size_t index = get_index(brow * mb, bcol * nb);
+    return &h_array[index];
+  }
+
+  // Fill with func(Matrix*,row, col)
+  void fill(T (*func)(matrix<T>*, int, int))
+  {
+    // Fill blocks by blocks
+    for (size_t colb = 0; colb < nt; colb++)
+    {
+      for (size_t rowb = 0; rowb < mt; rowb++)
+      {
+        T* addr_h = get_block_h(rowb, colb);
+#ifdef TILED
+        // tiles are stored contiguously
+        int ld = mb;
+#else
+        int ld = m;
+#endif
+
+        for (size_t lrow = 0; lrow < mb; lrow++)
+        {
+          for (size_t lcol = 0; lcol < nb; lcol++)
+          {
+            size_t row = lrow + rowb * mb;
+            size_t col = lcol + colb * nb;
+
+            T val = func(this, row, col);
+
+            addr_h[lrow + lcol * ld] = val;
+          }
+        }
+      }
+    }
+  }
+
+  T* h_array;
+  size_t m; // nrows
+  size_t n; // ncols
+
+  size_t mb; // block size (rows)
+  size_t nb; // block size (cols)
+
+  size_t mt; // numter of column blocks
+  size_t nt; // numter of row blocks
+
+  // abstract data handles
+  std::vector<logical_data<slice<T, 2>>> handles;
+
+  const char* symbol;
+
+  // for the mapping
+  int ndevs;
+  int grid_p, grid_q;
+};
+
+template <typename Ctx>
+void DGEMM(
+  Ctx& ctx,
+  cublasOperation_t transa,
+  cublasOperation_t transb,
+  double alpha,
+  matrix<double>& A,
+  int A_row,
+  int A_col,
+  matrix<double>& B,
+  int B_row,
+  int B_col,
+  double beta,
+  matrix<double>& C,
+  int C_row,
+  int C_col)
+{
+  auto dev = exec_place::device(C.get_preferred_devid(C_row, C_col));
+
+  auto t = ctx.task(
+    dev, A.get_handle(A_row, A_col).read(), B.get_handle(B_row, B_col).read(), C.get_handle(C_row, C_col).rw());
+  t.set_symbol("DGEMM");
+
+  t->*[&](cudaStream_t stream, auto tA, auto tB, auto tC) {
+    cuda_safe_call(cublasSetStream(get_cublas_handle(), stream));
+    int k = tA.extent(transa == CUBLAS_OP_N ? 1 : 0);
+    cuda_safe_call(cublasDgemm(
+      get_cublas_handle(),
+      transa,
+      transb,
+      tC.extent(0),
+      tC.extent(1),
+      k,
+      &alpha,
+      tA.data_handle(),
+      tA.stride(1),
+      tB.data_handle(),
+      tB.stride(1),
+      &beta,
+      tC.data_handle(),
+      tC.stride(1)));
+  };
+}
+
+template <typename Ctx>
+void PDGEMM(Ctx& ctx,
+            cublasOperation_t transa,
+            cublasOperation_t transb,
+            double alpha,
+            matrix<double>& A,
+            matrix<double>& B,
+            double beta,
+            matrix<double>& C)
+{
+  for (size_t m = 0; m < C.mt; m++)
+  {
+    for (size_t n = 0; n < C.nt; n++)
+    {
+      //=========================================
+      // alpha*A*B does not contribute; scale C
+      //=========================================
+      int inner_k = transa == CUBLAS_OP_N ? A.n : A.m;
+      if (alpha == 0.0 || inner_k == 0)
+      {
+        DGEMM(ctx, transa, transb, alpha, A, 0, 0, B, 0, 0, beta, C, m, n);
+      }
+      else if (transa == CUBLAS_OP_N)
+      {
+        //================================
+        // CUBLAS_OP_N / CUBLAS_OP_N
+        //================================
+        if (transb == CUBLAS_OP_N)
+        {
+          assert(A.nt == B.mt);
+          for (size_t k = 0; k < A.nt; k++)
+          {
+            double zbeta = k == 0 ? beta : 1.0;
+            DGEMM(ctx, transa, transb, alpha, A, m, k, B, k, n, zbeta, C, m, n);
+          }
+        }
+        //=====================================
+        // CUBLAS_OP_N / CUBLAS_OP_T
+        //=====================================
+        else
+        {
+          for (size_t k = 0; k < A.nt; k++)
+          {
+            double zbeta = k == 0 ? beta : 1.0;
+            DGEMM(ctx, transa, transb, alpha, A, m, k, B, n, k, zbeta, C, m, n);
+          }
+        }
+      }
+      else
+      {
+        //=====================================
+        // CUBLAS_OP_T / CUBLAS_OP_N
+        //=====================================
+        if (transb == CUBLAS_OP_N)
+        {
+          for (size_t k = 0; k < A.mt; k++)
+          {
+            double zbeta = k == 0 ? beta : 1.0;
+            DGEMM(ctx, transa, transb, alpha, A, k, m, B, k, n, zbeta, C, m, n);
+          }
+        }
+        //==========================================
+        // CUBLAS_OP_T / CUBLAS_OP_T
+        //==========================================
+        else
+        {
+          for (size_t k = 0; k < A.mt; k++)
+          {
+            double zbeta = k == 0 ? beta : 1.0;
+            DGEMM(ctx, transa, transb, alpha, A, k, m, B, n, k, zbeta, C, m, n);
+          }
+        }
+      }
+    }
+  }
+}
+
+double hilbert(matrix<double>* mat, int row, int col)
+{
+  return 1.0 / (col + row + 1.0) + 2.0 * mat->n * (col == row);
+}
+
+template <typename Ctx>
+void run(size_t N, size_t NB)
+{
+  /* This is the CUDASTF context */
+  Ctx ctx;
+
+  matrix<double> A(ctx, N, N, NB, NB, "A");
+  matrix<double> B(ctx, N, N, NB, NB, "B");
+  matrix<double> C(ctx, N, N, NB, NB, "C");
+
+  // (Hilbert matrix + 2*N*Id) to have a diagonal dominant matrix
+  A.fill(hilbert);
+  B.fill(hilbert);
+  C.fill(hilbert);
+
+  PDGEMM(ctx, CUBLAS_OP_N, CUBLAS_OP_N, 1.0, A, B, -2.0, C);
+
+  ctx.finalize();
+}
+
+int main(int argc, char** argv)
+{
+  size_t N  = 1024;
+  size_t NB = 128;
+
+  if (argc > 1)
+  {
+    N = atoi(argv[1]);
+  }
+
+  if (argc > 2)
+  {
+    NB = atoi(argv[2]);
+  }
+
+  assert(N % NB == 0);
+
+  run<stream_ctx>(N, NB);
+  run<graph_ctx>(N, NB);
+}

cccl_upstream/cudax/test/stf/gnu/07-cholesky.cpp

@@ -0,0 +1,718 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ *
+ * @brief This example implements a Cholesky decomposition over multiple devices using CUBLAS and CUSOLVER
+ *
+ * It also illustrates how we can use CUDASTF to allocate temporary data for CUSOLVER in CUDASTF tasks
+ */
+
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+#include <cuda/experimental/__stf/utility/nvtx.cuh>
+
+#include <iostream>
+
+#define TILED
+
+using namespace cuda::experimental::stf;
+
+// Global for the sake of simplicity !
+stream_ctx ctx;
+
+/* Get a CUBLAS handle valid on the current execution place, or initialize it lazily */
+cublasHandle_t& get_cublas_handle(const exec_place& ep = exec_place::current_device())
+{
+  static std::unordered_map<exec_place, cublasHandle_t, hash<exec_place>> cublas_handles;
+  auto& result = cublas_handles[ep];
+  if (result == cublasHandle_t())
+  { // not found, default value inserted
+    // Lazy initialization, and save the handle for future use
+    cuda_safe_call(cublasCreate(&result));
+  }
+  return result;
+}
+
+/* Get a CUSOLVER handle valid on the current execution place, or initialize it lazily */
+cusolverDnHandle_t& get_cusolver_handle(const exec_place& ep = exec_place::current_device())
+{
+  static std::unordered_map<exec_place, cusolverDnHandle_t, hash<exec_place>> cusolver_handles;
+  auto& result = cusolver_handles[ep];
+  if (result == cusolverDnHandle_t())
+  { // not found, default value inserted
+    // Lazy initialization, and save the handle for future use
+    cuda_safe_call(cusolverDnCreate(&result));
+  }
+  return result;
+}
+
+template <typename T>
+class matrix
+{
+public:
+  matrix(int NROWS, int NCOLS, int BLOCKSIZE_ROWS, int BLOCKSIZE_COLS, bool is_sym, const char* _symbol = "matrix")
+  {
+    symbol = _symbol;
+
+    sym_matrix = is_sym;
+
+    m  = NROWS;
+    mb = BLOCKSIZE_ROWS;
+
+    n  = NCOLS;
+    nb = BLOCKSIZE_COLS;
+
+    assert(m % mb == 0);
+    assert(n % nb == 0);
+
+    // cuda_safe_call(cudaMallocHost(&h_array, m*n*sizeof(T)));
+    // fprintf(stderr, "Allocating %ld x %ld x %ld = %ld bytes (%f GB) on host for %s\n", m, n, sizeof(T), s,
+    //        s / (1024.0 * 1024.0 * 1024.0), _symbol);
+    h_array.resize(m * n);
+    cuda_safe_call(cudaHostRegister(&h_array[0], h_array.size() * sizeof(T), cudaHostRegisterPortable));
+
+    // Compute the number of blocks
+    mt = m / mb;
+    nt = n / nb;
+
+    handles.resize(mt * nt);
+
+    for (size_t colb = 0; colb < nt; colb++)
+    {
+      int low_rowb = sym_matrix ? colb : 0;
+      for (size_t rowb = low_rowb; rowb < mt; rowb++)
+      {
+        T* addr_h = get_block_h(rowb, colb);
+        auto& h   = handle(rowb, colb);
+
+#ifdef TILED
+        // tiles are stored contiguously
+        size_t ld = mb;
+#else
+        size_t ld = m;
+#endif
+        std::ignore = ld; // work around bug in compiler
+        h           = ctx.logical_data(make_slice(addr_h, std::tuple{mb, nb}, ld));
+        h.set_symbol(std::string(symbol) + "_" + std::to_string(rowb) + "_" + std::to_string(colb));
+      }
+    }
+
+    cuda_safe_call(cudaGetDeviceCount(&ndevs));
+    for (int a = 1; a * a <= ndevs; a++)
+    {
+      if (ndevs % a == 0)
+      {
+        grid_p = a;
+        grid_q = ndevs / a;
+      }
+    }
+
+    assert(grid_p * grid_q == ndevs);
+
+    // std::cout << "FOUND " << ndevs << " DEVICES "
+    //          << "p=" << grid_p << " q=" << grid_q << '\n';
+  }
+
+  int get_preferred_devid(int row, int col)
+  {
+    return (row % grid_p) + (col % grid_q) * grid_p;
+  }
+
+  auto& handle(int row, int col)
+  {
+    return handles[row + col * mt];
+  }
+
+  size_t get_index(size_t row, size_t col)
+  {
+#ifdef TILED
+    // Find which tile contains this element
+    int tile_row = row / mb;
+    int tile_col = col / nb;
+
+    size_t tile_size = mb * nb;
+
+    // Look for the index of the beginning of the tile
+    size_t tile_start = (tile_row + mt * tile_col) * tile_size;
+
+    // Offset within the tile
+    size_t offset = (row % mb) + (col % nb) * mb;
+
+    return tile_start + offset;
+#else
+    return row + col * m;
+#endif
+  }
+
+  T* get_block_h(int brow, int bcol)
+  {
+    size_t index = get_index(brow * mb, bcol * nb);
+    return &h_array[index];
+  }
+
+  // Fill with func(Matrix*,row, col)
+  template <typename Fun>
+  void fill(Fun&& fun)
+  {
+    // Fill blocks by blocks
+    for (size_t colb = 0; colb < nt; colb++)
+    {
+      int low_rowb = sym_matrix ? colb : 0;
+      for (size_t rowb = low_rowb; rowb < mt; rowb++)
+      {
+        // Each task fills a block
+        ctx.host_launch(handle(rowb, colb).write())->*[this, fun, rowb, colb](auto sA) {
+          for (size_t lcol = 0; lcol < sA.extent(1); lcol++)
+          {
+            size_t col = lcol + colb * sA.extent(1);
+            for (size_t lrow = 0; lrow < sA.extent(0); lrow++)
+            {
+              size_t row     = lrow + rowb * sA.extent(0);
+              sA(lrow, lcol) = fun(*this, row, col);
+            }
+          }
+        };
+      }
+    }
+  }
+
+  std::vector<T> h_array;
+  size_t m; // nrows
+  size_t n; // ncols
+
+  // Is this a sym matrix ? (lower assumed)
+  bool sym_matrix;
+
+  size_t mb; // block size (rows)
+  size_t nb; // block size (cols)
+
+  size_t mt; // number of column blocks
+  size_t nt; // number of row blocks
+
+  // abstract data handles
+  std::vector<logical_data<slice<double, 2>>> handles;
+
+  const char* symbol;
+
+  // for the mapping
+  int ndevs;
+  int grid_p, grid_q;
+};
+
+void DPOTRF(cublasFillMode_t uplo, class matrix<double>& A, int A_row, int A_col)
+{
+  auto& Akk    = A.handle(A_row, A_col);
+  size_t m_akk = Akk.shape().extent(0);
+  // Note that the handle may be different from the actual handle...
+  int Lwork_expected;
+  cuda_safe_call(cusolverDnDpotrf_bufferSize(get_cusolver_handle(), uplo, m_akk, nullptr, 0, &Lwork_expected));
+
+  auto potrf_buffer = ctx.logical_data<double>(Lwork_expected);
+  auto devInfo      = ctx.logical_data(shape_of<slice<int>>(1));
+
+  auto t = ctx.task(Akk.rw(), potrf_buffer.write(), devInfo.write());
+  t.set_symbol("DPOTRF");
+  t->*[&](cudaStream_t s, auto sAkk, auto buffer, auto info) {
+    auto& h = get_cusolver_handle();
+    cuda_safe_call(cusolverDnSetStream(h, s));
+
+    cuda_safe_call(cusolverDnDpotrf(
+      h,
+      uplo,
+      sAkk.extent(0),
+      sAkk.data_handle(),
+      sAkk.stride(1),
+      buffer.data_handle(),
+      buffer.extent(0),
+      info.data_handle()));
+  };
+}
+
+void DGEMM(
+  cublasOperation_t transa,
+  cublasOperation_t transb,
+  double alpha,
+  class matrix<double>& A,
+  int A_row,
+  int A_col,
+  class matrix<double>& B,
+  int B_row,
+  int B_col,
+  double beta,
+  class matrix<double>& C,
+  int C_row,
+  int C_col)
+{
+  auto t = ctx.task(A.handle(A_row, A_col).read(), B.handle(B_row, B_col).read(), C.handle(C_row, C_col).rw());
+  t.set_symbol("DGEMM");
+  t->*[&](cudaStream_t s, auto sA, auto sB, auto sC) {
+    auto& h = get_cublas_handle();
+    cuda_safe_call(cublasSetStream(h, s));
+
+    auto k = (transa == CUBLAS_OP_N) ? sA.extent(1) : sA.extent(0);
+    cuda_safe_call(cublasDgemm(
+      h,
+      transa,
+      transb,
+      sC.extent(0),
+      sC.extent(1),
+      k,
+      &alpha,
+      sA.data_handle(),
+      sA.stride(1),
+      sB.data_handle(),
+      sB.stride(1),
+      &beta,
+      sC.data_handle(),
+      sC.stride(1)));
+  };
+}
+
+void DSYRK(
+  cublasFillMode_t uplo,
+  cublasOperation_t trans,
+  double alpha,
+  class matrix<double>& A,
+  int A_row,
+  int A_col,
+  double beta,
+  class matrix<double>& C,
+  int C_row,
+  int C_col)
+{
+  auto t = ctx.task(A.handle(A_row, A_col).read(), C.handle(C_row, C_col).rw());
+  t.set_symbol("DSYRK");
+  t->*[&](cudaStream_t s, auto sA, auto sC) {
+    auto& h = get_cublas_handle();
+    cuda_safe_call(cublasSetStream(h, s));
+
+    // number of rows of matrix op(A) and C
+    auto n = sC.extent(0);
+
+    // number of columns of matrix op(A)
+    auto k = (trans == CUBLAS_OP_N) ? sA.extent(1) : sA.extent(0);
+
+    cuda_safe_call(
+      cublasDsyrk(h, uplo, trans, n, k, &alpha, sA.data_handle(), sA.stride(1), &beta, sC.data_handle(), sC.stride(1)));
+  };
+}
+
+void DTRSM(
+  cublasSideMode_t side,
+  cublasFillMode_t uplo,
+  cublasOperation_t transa,
+  cublasDiagType_t diag,
+  double alpha,
+  class matrix<double>& A,
+  int A_row,
+  int A_col,
+  class matrix<double>& B,
+  int B_row,
+  int B_col)
+{
+  auto t = ctx.task(A.handle(A_row, A_col).read(), B.handle(B_row, B_col).rw());
+  t.set_symbol("DTRSM");
+  t->*[&](cudaStream_t s, auto sA, auto sB) {
+    auto& h = get_cublas_handle();
+    cuda_safe_call(cublasSetStream(h, s));
+
+    cuda_safe_call(cublasDtrsm(
+      h,
+      side,
+      uplo,
+      transa,
+      diag,
+      sB.extent(0),
+      sB.extent(1),
+      &alpha,
+      sA.data_handle(),
+      sA.stride(1),
+      sB.data_handle(),
+      sB.stride(1)));
+  };
+}
+
+void PDNRM2_HOST(matrix<double>* A, double* result)
+{
+#ifdef HAVE_DOT
+  reserved::dot::set_current_color("red");
+#endif
+
+  for (size_t rowb = 0; rowb < A->mt; rowb++)
+  {
+    for (size_t colb = 0; colb < A->nt; colb++)
+    {
+      ctx.host_launch(A->handle(rowb, colb).read())->*[=](auto sA) {
+        double res2 = 0.0;
+        for (size_t col = 0; col < sA.extent(1); col++)
+        {
+          for (size_t row = 0; row < sA.extent(0); row++)
+          {
+            double v = sA(row, col);
+            res2 += v * v;
+          }
+        }
+        *result += res2;
+      };
+    }
+  }
+}
+
+void PDPOTRF(matrix<double>& A)
+{
+  nvtx_range r("PDPOTRF");
+
+#ifdef HAVE_DOT
+  reserved::dot::set_current_color("yellow");
+#endif
+
+  assert(A.m == A.n);
+  assert(A.mt == A.nt);
+
+  int NBLOCKS = A.mt;
+  assert(A.mb == A.nb);
+
+  cuda_safe_call(cudaSetDevice(0));
+
+  for (int K = 0; K < NBLOCKS; K++)
+  {
+    cuda_safe_call(cudaSetDevice(A.get_preferred_devid(K, K)));
+    DPOTRF(CUBLAS_FILL_MODE_LOWER, A, K, K);
+
+    for (int row = K + 1; row < NBLOCKS; row++)
+    {
+      cuda_safe_call(cudaSetDevice(A.get_preferred_devid(row, K)));
+      DTRSM(CUBLAS_SIDE_RIGHT, CUBLAS_FILL_MODE_LOWER, CUBLAS_OP_T, CUBLAS_DIAG_NON_UNIT, 1.0, A, K, K, A, row, K);
+
+      for (int col = K + 1; col < row; col++)
+      {
+        cuda_safe_call(cudaSetDevice(A.get_preferred_devid(row, col)));
+        DGEMM(CUBLAS_OP_N, CUBLAS_OP_T, -1.0, A, row, K, A, col, K, 1.0, A, row, col);
+      }
+
+      cuda_safe_call(cudaSetDevice(A.get_preferred_devid(row, row)));
+      DSYRK(CUBLAS_FILL_MODE_LOWER, CUBLAS_OP_N, -1.0, A, row, K, 1.0, A, row, row);
+    }
+  }
+  cuda_safe_call(cudaSetDevice(0));
+}
+
+// Algorithm from PLASMA
+void PDTRSM(cublasSideMode_t side,
+            cublasFillMode_t uplo,
+            cublasOperation_t trans,
+            cublasDiagType_t diag,
+            double alpha,
+            class matrix<double>& A,
+            class matrix<double>& B)
+{
+  nvtx_range r("PDTRSM");
+
+  //    std::cout << "[PDTRSM] START B MT " << B.mt << " NT " << B.nt << '\n';
+
+  if (side == CUBLAS_SIDE_LEFT)
+  {
+    if (uplo == CUBLAS_FILL_MODE_UPPER)
+    {
+      // TODO
+      assert(0);
+      abort();
+    }
+    else
+    {
+      //===========================================
+      // CUBLAS_SIDE_LEFT / CUBLAS_FILL_MODE_LOWER / CUBLAS_OP_N
+      //===========================================
+      if (trans == CUBLAS_OP_N)
+      {
+        for (size_t k = 0; k < B.mt; k++)
+        {
+          double lalpha = k == 0 ? alpha : 1.0;
+          for (size_t n = 0; n < B.nt; n++)
+          {
+            cuda_safe_call(cudaSetDevice(A.get_preferred_devid(k, k)));
+            DTRSM(side, uplo, trans, diag, lalpha, A, k, k, B, k, n);
+          }
+          for (size_t m = k + 1; m < B.mt; m++)
+          {
+            for (size_t n = 0; n < B.nt; n++)
+            {
+              cuda_safe_call(cudaSetDevice(A.get_preferred_devid(m, k)));
+              DGEMM(CUBLAS_OP_N, CUBLAS_OP_N, -1.0, A, m, k, B, k, n, lalpha, B, m, n);
+            }
+          }
+        }
+      }
+      //================================================
+      // CUBLAS_SIDE_LEFT / CUBLAS_FILL_MODE_LOWER / CUBLAS_OP_[C|T]
+      //================================================
+      else
+      {
+        for (size_t k = 0; k < B.mt; k++)
+        {
+          double lalpha = k == 0 ? alpha : 1.0;
+          for (size_t n = 0; n < B.nt; n++)
+          {
+            cuda_safe_call(cudaSetDevice(A.get_preferred_devid(B.mt - k - 1, B.mt - k - 1)));
+            DTRSM(side, uplo, trans, diag, lalpha, A, B.mt - k - 1, B.mt - k - 1, B, B.mt - k - 1, n);
+          }
+          for (size_t m = k + 1; m < B.mt; m++)
+          {
+            for (size_t n = 0; n < B.nt; n++)
+            {
+              cuda_safe_call(cudaSetDevice(A.get_preferred_devid(B.mt - k - 1, B.mt - 1 - m)));
+              DGEMM(
+                trans, CUBLAS_OP_N, -1.0, A, B.mt - k - 1, B.mt - 1 - m, B, B.mt - k - 1, n, lalpha, B, B.mt - 1 - m, n);
+            }
+          }
+        }
+      }
+    }
+  }
+  else
+  {
+    // TODO
+    abort();
+  }
+  cuda_safe_call(cudaSetDevice(0));
+  //    std::cout << "[PDTRSM] END" << '\n';
+}
+
+void PDPOTRS(matrix<double>& A, class matrix<double>& B, cublasFillMode_t uplo)
+{
+  nvtx_range r("PDPOTRS");
+
+#ifdef HAVE_DOT
+  reserved::dot::set_current_color("green");
+#endif
+
+  //    std::cout << "[PDPOTRS] START" << '\n';
+  // Call the parallel functions.
+  PDTRSM(
+    CUBLAS_SIDE_LEFT, uplo, uplo == CUBLAS_FILL_MODE_UPPER ? CUBLAS_OP_T : CUBLAS_OP_N, CUBLAS_DIAG_NON_UNIT, 1.0, A, B);
+
+#ifdef HAVE_DOT
+  reserved::dot::set_current_color("darkgreen");
+#endif
+
+  PDTRSM(
+    CUBLAS_SIDE_LEFT, uplo, uplo == CUBLAS_FILL_MODE_UPPER ? CUBLAS_OP_N : CUBLAS_OP_T, CUBLAS_DIAG_NON_UNIT, 1.0, A, B);
+  //    std::cout << "[PDPOTRS] END" << '\n';
+}
+
+/*****************************************************************************
+ * Parallel tile matrix-matrix
+ *multiplication.
+ * @see plasma_omp_dgemm
+ ******************************************************************************/
+void PDGEMM(cublasOperation_t transa,
+            cublasOperation_t transb,
+            double alpha,
+            class matrix<double>& A,
+            class matrix<double>& B,
+            double beta,
+            class matrix<double>& C)
+{
+  nvtx_range r("PDGEMM");
+
+#ifdef HAVE_DOT
+  reserved::dot::set_current_color("blue");
+#endif
+
+  for (size_t m = 0; m < C.mt; m++)
+  {
+    for (size_t n = 0; n < C.nt; n++)
+    {
+      //=========================================
+      // alpha*A*B does not contribute; scale C
+      //=========================================
+      int inner_k = transa == CUBLAS_OP_N ? A.n : A.m;
+      if (alpha == 0.0 || inner_k == 0)
+      {
+        DGEMM(transa, transb, alpha, A, 0, 0, B, 0, 0, beta, C, m, n);
+      }
+      else if (transa == CUBLAS_OP_N)
+      {
+        //================================
+        // CUBLAS_OP_N / CUBLAS_OP_N
+        //================================
+        if (transb == CUBLAS_OP_N)
+        {
+          for (size_t k = 0; k < A.nt; k++)
+          {
+            double zbeta = k == 0 ? beta : 1.0;
+            DGEMM(transa, transb, alpha, A, m, k, B, k, n, zbeta, C, m, n);
+          }
+        }
+        //=====================================
+        // CUBLAS_OP_N / CUBLAS_OP_T
+        //=====================================
+        else
+        {
+          for (size_t k = 0; k < A.nt; k++)
+          {
+            double zbeta = k == 0 ? beta : 1.0;
+            DGEMM(transa, transb, alpha, A, m, k, B, n, k, zbeta, C, m, n);
+          }
+        }
+      }
+      else
+      {
+        //=====================================
+        // CUBLAS_OP_T / CUBLAS_OP_N
+        //=====================================
+        if (transb == CUBLAS_OP_N)
+        {
+          for (size_t k = 0; k < A.mt; k++)
+          {
+            double zbeta = k == 0 ? beta : 1.0;
+            DGEMM(transa, transb, alpha, A, k, m, B, k, n, zbeta, C, m, n);
+          }
+        }
+        //==========================================
+        // CUBLAS_OP_T / CUBLAS_OP_T
+        //==========================================
+        else
+        {
+          for (size_t k = 0; k < A.mt; k++)
+          {
+            double zbeta = k == 0 ? beta : 1.0;
+            DGEMM(transa, transb, alpha, A, k, m, B, n, k, zbeta, C, m, n);
+          }
+        }
+      }
+    }
+  }
+}
+
+int main(int argc, char** argv)
+{
+  int N  = 1024;
+  int NB = 128;
+
+  if (argc > 1)
+  {
+    N = atoi(argv[1]);
+  }
+
+  if (argc > 2)
+  {
+    NB = atoi(argv[2]);
+  }
+
+  int check_result = 1;
+  if (getenv("CHECK_RESULT"))
+  {
+    check_result = atoi(getenv("CHECK_RESULT"));
+  }
+
+  assert(N % NB == 0);
+
+  // Set up CUBLAS and CUSOLVER
+  int ndevs;
+  cuda_safe_call(cudaGetDeviceCount(&ndevs));
+
+  cuda_safe_call(cudaSetDevice(0));
+
+  matrix<double> A(N, N, NB, NB, true, "A");
+  matrix<double> Aref(N, N, NB, NB, false, "Aref");
+
+  // (Hilbert matrix + 2*N*Id) to have a diagonal dominant matrix
+  auto hilbert = [](matrix<double>& mat, int row, int col) {
+    return 1.0 / (col + row + 1.0) + 2.0 * mat.n * (col == row);
+  };
+
+  if (check_result)
+  {
+    Aref.fill(hilbert);
+  }
+
+  A.fill(hilbert);
+
+  /* Right-hand side */
+  matrix<double> B_potrs(N, 1, NB, 1, false, "B");
+  matrix<double> Bref_potrs(N, 1, NB, 1, false, "Bref");
+
+  if (check_result)
+  {
+    auto rhs_vals = [](matrix<double>&, int row, int /*col*/) {
+      return 1.0 * (row + 1);
+    };
+    B_potrs.fill(rhs_vals);
+    Bref_potrs.fill(rhs_vals);
+  }
+
+  //    // Compute ||Bref||
+  double Bref_nrm2 = 0.0;
+  double res_nrm2  = 0.0;
+
+  if (check_result)
+  {
+    PDNRM2_HOST(&Bref_potrs, &Bref_nrm2);
+  }
+
+  cudaEvent_t startEvent_pdpotrf, stopEvent_pdpotrf;
+  float milliseconds_pdpotrf = 0;
+
+  //    for (size_t row = 0; row < A.mt; row++)
+  //    {
+  //        for (size_t col = 0; col <= row; col++)
+  //        {
+  //            cuda_safe_call(cudaSetDevice(A.get_preferred_devid(row, col)));
+  //            NOOP(A, row, col);
+  //        }
+  //    }
+
+  cuda_safe_call(cudaEventCreate(&startEvent_pdpotrf));
+  cuda_safe_call(cudaEventCreate(&stopEvent_pdpotrf));
+
+  cuda_safe_call(cudaEventRecord(startEvent_pdpotrf, ctx.fence()));
+
+  PDPOTRF(A);
+
+  cuda_safe_call(cudaEventRecord(stopEvent_pdpotrf, ctx.fence()));
+
+  /*
+   *  POTRS
+   */
+
+  if (check_result)
+  {
+    // Solve AX = B and put the result in B
+    PDPOTRS(A, B_potrs, CUBLAS_FILL_MODE_LOWER);
+
+    // Compute (AX - B)
+    // Bref = (Aref*B - Bref)
+    PDGEMM(CUBLAS_OP_N, CUBLAS_OP_N, 1.0, Aref, B_potrs, -1.0, Bref_potrs);
+
+    // Compute ||AX - B|| = ||Bref||
+    PDNRM2_HOST(&Bref_potrs, &res_nrm2);
+  }
+
+  ctx.finalize();
+
+  cuda_safe_call(cudaEventElapsedTime(&milliseconds_pdpotrf, startEvent_pdpotrf, stopEvent_pdpotrf));
+
+  double gflops_pdpotrf = 1.0 / 3.0 * ((double) N * (double) N * (double) N) / (1000000000.0);
+  std::cout << "[PDPOTRF] ELAPSED: " << milliseconds_pdpotrf
+            << " ms, GFLOPS: " << gflops_pdpotrf / (milliseconds_pdpotrf / 1000.0) << '\n';
+
+  if (check_result)
+  {
+    if (const auto residual = sqrt(res_nrm2) / sqrt(Bref_nrm2); residual >= 0.01)
+    {
+      std::cerr << "[POTRS] ||AX - B|| : " << sqrt(res_nrm2) << '\n';
+      std::cerr << "[POTRS] ||B|| : " << sqrt(Bref_nrm2) << '\n';
+      std::cerr << "[POTRS] RESIDUAL (||AX - B||/||B||) : " << residual << '\n';
+      assert(!"Algorithm did not converge.");
+    }
+  }
+}

cccl_upstream/cudax/test/stf/gnu/include_only.cpp

@@ -0,0 +1,23 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+int main()
+{
+  graph_ctx ctx;
+  ctx.finalize();
+
+  stream_ctx ctx2;
+  ctx2.finalize();
+}

cccl_upstream/cudax/test/stf/graph/concurrency_test.cu

@@ -0,0 +1,86 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
+
+#include <iostream>
+
+/*
+ * The goal of this test is to ensure that using read access modes actually
+ * results in concurrent tasks
+ */
+
+using namespace cuda::experimental::stf;
+
+static __global__ void cuda_sleep_kernel(long long int clock_cnt)
+{
+  long long int start_clock  = clock64();
+  long long int clock_offset = 0;
+  while (clock_offset < clock_cnt)
+  {
+    clock_offset = clock64() - start_clock;
+  }
+}
+
+int main(int argc, char** argv)
+{
+  int NTASKS = 256;
+  int ms     = 40;
+
+  if (argc > 1)
+  {
+    NTASKS = atoi(argv[1]);
+  }
+
+  if (argc > 2)
+  {
+    ms = atoi(argv[2]);
+  }
+
+  // cudaDevAttrClockRate: Peak clock frequency in kilohertz;
+  int clock_rate;
+  cuda_safe_call(cudaDeviceGetAttribute(&clock_rate, cudaDevAttrClockRate, 0));
+  long long int clock_cnt = (long long int) (ms * clock_rate);
+
+  graph_ctx ctx;
+
+  int dummy[1];
+  auto handle = ctx.logical_data(dummy);
+
+  ctx.task(handle.rw())->*[](cudaGraph_t graph, auto /*unused*/) {
+    cudaGraphNode_t n;
+    cuda_safe_call(cudaGraphAddEmptyNode(&n, graph, nullptr, 0));
+  };
+
+  for (int iter = 0; iter < 10; iter++)
+  {
+    for (int k = 0; k < NTASKS; k++)
+    {
+      ctx.task(handle.read())->*[&](cudaStream_t stream, auto /*unused*/) {
+        cuda_sleep_kernel<<<1, 1, 0, stream>>>(clock_cnt);
+      };
+    }
+
+    ctx.task(handle.rw())->*[&](cudaGraph_t graph, auto /*unused*/) {
+      cudaGraphNode_t n;
+      cuda_safe_call(cudaGraphAddEmptyNode(&n, graph, nullptr, 0));
+    };
+  }
+
+  ctx.submit();
+
+  if (argc > 3)
+  {
+    std::cout << "Generating DOT output in " << argv[3] << '\n';
+    ctx.print_to_dot(argv[3]);
+  }
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/graph/epoch.cu

@@ -0,0 +1,56 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief Test explicit uses of the API to change stage and create a sequence
+ *        of CUDA graphs
+ */
+
+#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+int main()
+{
+  graph_ctx ctx;
+
+  const size_t N     = 8;
+  const size_t NITER = 2;
+
+  double A[N];
+  for (size_t i = 0; i < N; i++)
+  {
+    A[i] = 1.0 * i;
+  }
+
+  auto lA = ctx.logical_data(A);
+
+  for (size_t k = 0; k < NITER; k++)
+  {
+    ctx.parallel_for(blocked_partition(), exec_place::current_device(), lA.shape(), lA.rw())
+        ->*[] __host__ __device__(size_t i, slice<double> A) { A(i) = cos(A(i)); };
+
+    ctx.change_stage();
+  }
+
+  ctx.finalize();
+
+  for (size_t i = 0; i < N; i++)
+  {
+    double Ai_ref = 1.0 * i;
+    for (size_t k = 0; k < NITER; k++)
+    {
+      Ai_ref = cos(Ai_ref);
+    }
+
+    EXPECT(fabs(A[i] - Ai_ref) < 0.01);
+  }
+}

cccl_upstream/cudax/test/stf/graph/explicit_graph.cu

@@ -0,0 +1,57 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2025 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+//! @file
+//! @brief Add tasks to a user-provided graph
+
+#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+__global__ void dummy() {}
+
+int main()
+{
+  cudaGraph_t graph;
+  cudaGraphExec_t graphExec = NULL;
+  cudaStream_t stream;
+
+  cuda_safe_call(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
+
+  cuda_safe_call(cudaGraphCreate(&graph, 0));
+
+  graph_ctx ctx(graph);
+
+  auto lX = ctx.token();
+  auto lY = ctx.token();
+  auto lZ = ctx.token();
+
+  ctx.task(lX.write())->*[](cudaStream_t s) {
+    dummy<<<1, 1, 0, s>>>();
+  };
+
+  ctx.task(lX.read(), lY.write())->*[](cudaStream_t s) {
+    dummy<<<1, 1, 0, s>>>();
+  };
+
+  ctx.task(lX.read(), lZ.write())->*[](cudaStream_t s) {
+    dummy<<<1, 1, 0, s>>>();
+  };
+
+  ctx.task(lY.rw(), lZ.rw())->*[](cudaStream_t s) {
+    dummy<<<1, 1, 0, s>>>();
+  };
+
+  ctx.finalize_as_graph();
+
+  cuda_safe_call(cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0));
+  cuda_safe_call(cudaGraphLaunch(graphExec, stream));
+  cuda_safe_call(cudaStreamSynchronize(stream));
+}

cccl_upstream/cudax/test/stf/graph/explicit_graph_async.cu

@@ -0,0 +1,57 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2025 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+//! @file
+//! @brief Add tasks to a user-provided graph and launch on a user-provided stream.
+//!        Exercises graph_ctx(cudaGraph_t, cudaStream_t): finalize() submits the
+//!        graph on the given stream and does not block; the caller synchronizes.
+
+#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+__global__ void dummy() {}
+
+int main()
+{
+  cudaGraph_t graph;
+  cudaStream_t stream;
+
+  cuda_safe_call(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
+  cuda_safe_call(cudaGraphCreate(&graph, 0));
+
+  graph_ctx ctx(graph, stream);
+
+  auto lX = ctx.token();
+  auto lY = ctx.token();
+  auto lZ = ctx.token();
+
+  ctx.task(lX.write())->*[](cudaStream_t s) {
+    dummy<<<1, 1, 0, s>>>();
+  };
+
+  ctx.task(lX.read(), lY.write())->*[](cudaStream_t s) {
+    dummy<<<1, 1, 0, s>>>();
+  };
+
+  ctx.task(lX.read(), lZ.write())->*[](cudaStream_t s) {
+    dummy<<<1, 1, 0, s>>>();
+  };
+
+  ctx.task(lY.rw(), lZ.rw())->*[](cudaStream_t s) {
+    dummy<<<1, 1, 0, s>>>();
+  };
+
+  // Non-blocking: submits the graph on the user-provided stream
+  ctx.finalize();
+
+  cuda_safe_call(cudaStreamSynchronize(stream));
+  cuda_safe_call(cudaStreamDestroy(stream));
+}

cccl_upstream/cudax/test/stf/graph/explicit_graph_while-kernels.cu

@@ -0,0 +1,92 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2025 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+//! @file
+//! @brief Add tasks to a user-provided child graph from a while loop
+
+#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+#if _CCCL_CTK_AT_LEAST(12, 4)
+__global__ void dummy() {}
+
+__global__ void setHandle(cudaGraphConditionalHandle handle)
+{
+  static int count = 5;
+  cudaGraphSetConditional(handle, --count ? 1 : 0);
+}
+#endif // _CCCL_CTK_AT_LEAST(12, 4)
+
+int main()
+{
+#if _CCCL_CTK_BELOW(12, 4)
+  fprintf(stderr, "Waiving test: conditional nodes are only available since CUDA 12.4.\n");
+#else
+  cudaStream_t stream;
+
+  cuda_safe_call(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
+
+  cudaGraph_t graph;
+  cudaGraphNode_t conditionalNode;
+
+  cudaGraphCreate(&graph, 0);
+
+  cudaGraphConditionalHandle handle;
+  cudaGraphConditionalHandleCreate(&handle, graph, 1, cudaGraphCondAssignDefault);
+
+  cudaGraphNodeParams cParams = {};
+  cParams.type                = cudaGraphNodeTypeConditional;
+  cParams.conditional.handle  = handle;
+  cParams.conditional.type    = cudaGraphCondTypeWhile;
+  cParams.conditional.size    = 1;
+#  if _CCCL_CTK_AT_LEAST(13, 0)
+  cudaGraphAddNode(&conditionalNode, graph, nullptr, nullptr, 0, &cParams);
+#  else
+  cudaGraphAddNode(&conditionalNode, graph, nullptr, 0, &cParams);
+#  endif
+
+  cudaGraph_t bodyGraph = cParams.conditional.phGraph_out[0];
+
+  graph_ctx ctx(bodyGraph);
+
+  auto lX = ctx.token();
+  auto lY = ctx.token();
+  auto lZ = ctx.token();
+
+  ctx.cuda_kernel(lX.write())->*[]() {
+    return cuda_kernel_desc{dummy, 1, 1, 0};
+  };
+
+  ctx.cuda_kernel(lX.read(), lY.write())->*[]() {
+    return cuda_kernel_desc{dummy, 1, 1, 0};
+  };
+
+  ctx.cuda_kernel(lX.read(), lZ.write())->*[]() {
+    return cuda_kernel_desc{dummy, 1, 1, 0};
+  };
+
+  ctx.cuda_kernel(lY.rw(), lZ.rw())->*[]() {
+    return cuda_kernel_desc{dummy, 1, 1, 0};
+  };
+
+  ctx.cuda_kernel()->*[handle]() {
+    return cuda_kernel_desc{setHandle, 1, 1, 0, handle};
+  };
+
+  ctx.finalize_as_graph();
+
+  cudaGraphExec_t graphExec = NULL;
+  cuda_safe_call(cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0));
+  cuda_safe_call(cudaGraphLaunch(graphExec, stream));
+  cuda_safe_call(cudaStreamSynchronize(stream));
+  cuda_safe_call(cudaGraphDebugDotPrint(graph, "test-while.dot", cudaGraphDebugDotFlags(0)));
+#endif // !_CCCL_CTK_BELOW(12, 4)
+}

cccl_upstream/cudax/test/stf/graph/explicit_graph_while.cu

@@ -0,0 +1,92 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2025 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+//! @file
+//! @brief Add tasks to a user-provided child graph from a while loop
+
+#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+#if _CCCL_CTK_AT_LEAST(12, 4)
+__global__ void dummy() {}
+
+__global__ void setHandle(cudaGraphConditionalHandle handle)
+{
+  static int count = 5;
+  cudaGraphSetConditional(handle, --count ? 1 : 0);
+}
+#endif // _CCCL_CTK_AT_LEAST(12, 4)
+
+int main()
+{
+#if _CCCL_CTK_BELOW(12, 4)
+  fprintf(stderr, "Waiving test: conditional nodes are only available since CUDA 12.4.\n");
+#else
+  cudaStream_t stream;
+
+  cuda_safe_call(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
+
+  cudaGraph_t graph;
+  cudaGraphNode_t conditionalNode;
+
+  cudaGraphCreate(&graph, 0);
+
+  cudaGraphConditionalHandle handle;
+  cudaGraphConditionalHandleCreate(&handle, graph, 1, cudaGraphCondAssignDefault);
+
+  cudaGraphNodeParams cParams = {};
+  cParams.type                = cudaGraphNodeTypeConditional;
+  cParams.conditional.handle  = handle;
+  cParams.conditional.type    = cudaGraphCondTypeWhile;
+  cParams.conditional.size    = 1;
+#  if _CCCL_CTK_AT_LEAST(13, 0)
+  cudaGraphAddNode(&conditionalNode, graph, nullptr, nullptr, 0, &cParams);
+#  else
+  cudaGraphAddNode(&conditionalNode, graph, nullptr, 0, &cParams);
+#  endif
+
+  cudaGraph_t bodyGraph = cParams.conditional.phGraph_out[0];
+
+  graph_ctx ctx(bodyGraph);
+
+  auto lX = ctx.token();
+  auto lY = ctx.token();
+  auto lZ = ctx.token();
+
+  ctx.task(lX.write())->*[](cudaStream_t s) {
+    dummy<<<1, 1, 0, s>>>();
+  };
+
+  ctx.task(lX.read(), lY.write())->*[](cudaStream_t s) {
+    dummy<<<1, 1, 0, s>>>();
+  };
+
+  ctx.task(lX.read(), lZ.write())->*[](cudaStream_t s) {
+    dummy<<<1, 1, 0, s>>>();
+  };
+
+  ctx.task(lY.rw(), lZ.rw())->*[](cudaStream_t s) {
+    dummy<<<1, 1, 0, s>>>();
+  };
+
+  ctx.task()->*[&handle](cudaStream_t s) {
+    setHandle<<<1, 1, 0, s>>>(handle);
+  };
+
+  ctx.finalize_as_graph();
+
+  cudaGraphExec_t graphExec = NULL;
+  cuda_safe_call(cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0));
+  cuda_safe_call(cudaGraphLaunch(graphExec, stream));
+  cuda_safe_call(cudaStreamSynchronize(stream));
+  cuda_safe_call(cudaGraphDebugDotPrint(graph, "test-while.dot", cudaGraphDebugDotFlags(0)));
+#endif // !_CCCL_CTK_BELOW(12, 4)
+}

cccl_upstream/cudax/test/stf/graph/freeze_for_graph.cu

@@ -0,0 +1,74 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ *
+ * @brief Ensure temporary data are destroyed
+ *
+ */
+
+#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+int X0(int i)
+{
+  return 17 * i + 45;
+}
+
+__global__ void dummy() {}
+
+int main()
+{
+  stream_ctx ctx;
+  const int N = 16;
+  int X[N];
+
+  for (int i = 0; i < N; i++)
+  {
+    X[i] = X0(i);
+  }
+
+  auto lX = ctx.logical_data(X);
+
+  auto fX = ctx.freeze(lX, access_mode::rw, data_place::current_device());
+
+  auto stream = ctx.pick_stream();
+
+  graph_ctx gctx(stream);
+
+  auto frozen_X = fX.get(data_place::current_device(), stream);
+  auto lX_alias = gctx.logical_data(frozen_X, data_place::current_device());
+
+  auto lY = gctx.logical_data(lX.shape());
+
+  gctx.parallel_for(lX.shape(), lX_alias.read(), lY.write())->*[] __device__(size_t i, auto x, auto y) {
+    y(i) = x(i);
+  };
+
+  gctx.parallel_for(lX.shape(), lX_alias.write(), lY.read())->*[] __device__(size_t i, auto x, auto y) {
+    x(i) = y(i) + 2;
+  };
+
+  gctx.finalize();
+
+  fX.unfreeze(stream);
+
+  ctx.host_launch(lX.read())->*[](auto x) {
+    for (int i = 0; i < static_cast<int>(x.size()); i++)
+    {
+      EXPECT(x(i) == X0(i) + 2);
+    }
+  };
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/graph/freeze_nested_graphs.cu

@@ -0,0 +1,96 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2025 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+//! \file
+//!
+//! \brief Freeze a logical data in a graph to use it in a child graph
+
+#include <cuda/experimental/stf.cuh>
+
+#include <vector>
+
+using namespace cuda::experimental::stf;
+
+int X0(int i)
+{
+  return 17 * i + 45;
+}
+
+__global__ void dummy() {}
+
+int main()
+{
+  const int N = 16;
+  int X[N];
+
+  for (int i = 0; i < N; i++)
+  {
+    X[i] = X0(i);
+  }
+
+  graph_ctx ctx;
+
+  auto lX = ctx.logical_data(X);
+
+  ctx.parallel_for(lX.shape(), lX.rw())->*[] __device__(size_t i, auto x) {
+    x(i) *= 3;
+  };
+
+  auto fX = ctx.freeze(lX, access_mode::rw, data_place::current_device());
+
+  // Create a graph that will later be inserted as a child graph once all input
+  // dependencies are known
+  cudaGraph_t sub_graph;
+  cuda_safe_call(cudaGraphCreate(&sub_graph, 0));
+
+  // Create a context based on this child graph
+  graph_ctx sub_ctx(sub_graph);
+
+  auto [frozen_X, fX_get_events] = fX.get(data_place::current_device());
+
+  auto lX_alias = sub_ctx.logical_data(frozen_X, data_place::current_device());
+
+  // XXX we need an adapter to allocate data from the upper context
+  //  auto lY = sub_ctx.logical_data(lX.shape());
+  //  sub_ctx.parallel_for(lX.shape(), lX_alias.read(), lY.write())->*[] __device__(size_t i, auto x, auto y) {
+  //    y(i) = x(i);
+  //  };
+
+  sub_ctx.parallel_for(lX.shape(), lX_alias.rw())->*[] __device__(size_t i, auto x) {
+    x(i) = x(i) + 2;
+  };
+
+  sub_ctx.finalize_as_graph();
+
+  // The child graph depends on the events to get the frozen data
+  ::std::vector<cudaGraphNode_t> fX_ready_nodes = reserved::join_with_graph_nodes(ctx, fX_get_events, ctx.stage());
+
+  // Add the child graph as a node that depends on the frozen data being ready
+  cudaGraphNode_t child_graph_node;
+  cuda_safe_call(cudaGraphAddChildGraphNode(
+    &child_graph_node, ctx.get_graph(), fX_ready_nodes.data(), fX_ready_nodes.size(), sub_ctx.get_graph()));
+
+  // Create an event that signals when the child graph completes
+  event_list child_graph_event;
+  reserved::fork_from_graph_node(
+    ctx, child_graph_node, ctx.get_graph(), ctx.stage(), child_graph_event, "child graph done");
+
+  // Unfreeze the data after the child graph completes
+  fX.unfreeze(child_graph_event);
+
+  ctx.host_launch(lX.read())->*[](auto x) {
+    for (int i = 0; i < static_cast<int>(x.size()); i++)
+    {
+      EXPECT(x(i) == 3 * X0(i) + 2);
+    }
+  };
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/graph/freeze_nested_graphs_allocator.cu

@@ -0,0 +1,96 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2025 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+//! \file
+//!
+//! \brief Freeze a logical data in a graph to use it in a child graph
+
+#include <cuda/experimental/stf.cuh>
+
+#include <vector>
+
+using namespace cuda::experimental::stf;
+
+int X0(int i)
+{
+  return 17 * i + 45;
+}
+
+__global__ void dummy() {}
+
+int main()
+{
+  const int N = 16;
+  int X[N];
+
+  for (int i = 0; i < N; i++)
+  {
+    X[i] = X0(i);
+  }
+
+  graph_ctx ctx;
+
+  auto lX = ctx.logical_data(X);
+
+  ctx.parallel_for(lX.shape(), lX.rw())->*[] __device__(size_t i, auto x) {
+    x(i) *= 3;
+  };
+
+  auto fX = ctx.freeze(lX, access_mode::rw, data_place::current_device());
+
+  // Create a graph that will later be inserted as a child graph once all input
+  // dependencies are known
+  cudaGraph_t sub_graph;
+  cuda_safe_call(cudaGraphCreate(&sub_graph, 0));
+
+  // Create a context based on this child graph
+  graph_ctx sub_ctx(sub_graph);
+
+  auto [frozen_X, fX_get_events] = fX.get(data_place::current_device());
+
+  auto lX_alias = sub_ctx.logical_data(frozen_X, data_place::current_device());
+
+  // XXX we need an adapter to allocate data from the upper context
+  //  auto lY = sub_ctx.logical_data(lX.shape());
+  //  sub_ctx.parallel_for(lX.shape(), lX_alias.read(), lY.write())->*[] __device__(size_t i, auto x, auto y) {
+  //    y(i) = x(i);
+  //  };
+
+  sub_ctx.parallel_for(lX.shape(), lX_alias.rw())->*[] __device__(size_t i, auto x) {
+    x(i) = x(i) + 2;
+  };
+
+  sub_ctx.finalize_as_graph();
+
+  // The child graph depends on the events to get the frozen data
+  ::std::vector<cudaGraphNode_t> fX_ready_nodes = reserved::join_with_graph_nodes(ctx, fX_get_events, ctx.stage());
+
+  // Add the child graph as a node that depends on the frozen data being ready
+  cudaGraphNode_t child_graph_node;
+  cuda_safe_call(cudaGraphAddChildGraphNode(
+    &child_graph_node, ctx.get_graph(), fX_ready_nodes.data(), fX_ready_nodes.size(), sub_ctx.get_graph()));
+
+  // Create an event that signals when the child graph completes
+  event_list child_graph_event;
+  reserved::fork_from_graph_node(
+    ctx, child_graph_node, ctx.get_graph(), ctx.stage(), child_graph_event, "child graph done");
+
+  // Unfreeze the data after the child graph completes
+  fX.unfreeze(child_graph_event);
+
+  ctx.host_launch(lX.read())->*[](auto x) {
+    for (int i = 0; i < static_cast<int>(x.size()); i++)
+    {
+      EXPECT(x(i) == 3 * X0(i) + 2);
+    }
+  };
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/graph/freeze_while_graphs.cu

@@ -0,0 +1,127 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2025 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+//! \file
+//!
+//! \brief Freeze a logical data in a graph to use it in the body of a "while" graph node
+
+#include <cuda/experimental/stf.cuh>
+
+#include <vector>
+
+using namespace cuda::experimental::stf;
+
+#if _CCCL_CTK_AT_LEAST(12, 4)
+int X0(int i)
+{
+  return 17 * i + 45;
+}
+
+__global__ void setHandle(cudaGraphConditionalHandle handle)
+{
+  static int count = 5;
+  cudaGraphSetConditional(handle, --count ? 1 : 0);
+}
+#endif // _CCCL_CTK_AT_LEAST(12, 4)
+
+int main()
+{
+#if _CCCL_CTK_BELOW(12, 4)
+  fprintf(stderr, "Waiving test: conditional nodes are only available since CUDA 12.4.\n");
+#else
+  const int N = 16;
+  int X[N];
+
+  for (int i = 0; i < N; i++)
+  {
+    X[i] = X0(i);
+  }
+
+  graph_ctx ctx;
+
+  auto lX = ctx.logical_data(X);
+
+  ctx.parallel_for(lX.shape(), lX.rw())->*[] __device__(size_t i, auto x) {
+    x(i) *= 3;
+  };
+
+  auto fX = ctx.freeze(lX, access_mode::rw, data_place::current_device());
+
+  cudaGraphConditionalHandle handle;
+  cudaGraphConditionalHandleCreate(&handle, ctx.get_graph(), 1, cudaGraphCondAssignDefault);
+
+  // Create a graph that will later be inserted as a child graph once all input
+  // dependencies are known
+  cudaGraph_t sub_graph;
+  cuda_safe_call(cudaGraphCreate(&sub_graph, 0));
+
+  // Create a context based on this child graph which is the body of the
+  graph_ctx sub_ctx(sub_graph);
+
+  auto [frozen_X, fX_get_events] = fX.get(data_place::current_device());
+
+  auto lX_alias = sub_ctx.logical_data(frozen_X, data_place::current_device());
+
+  sub_ctx.parallel_for(lX.shape(), lX_alias.rw())->*[] __device__(size_t i, auto x) {
+    x(i) = x(i) + 2;
+  };
+
+  // We want to repeat this a fixed number of times
+  sub_ctx.cuda_kernel()->*[handle]() {
+    return cuda_kernel_desc{setHandle, 1, 1, 0, handle};
+  };
+
+  sub_ctx.finalize_as_graph();
+
+  // We now create a conditional graph which depends on the same dependencies
+  // as the inner ctx. We then insert the body of the graph as a child graph of
+  // the conditional node because we cannot decide what graph is the body of
+  // the conditional node ourselves, and we cannot add input dependencies to
+  // the conditional node after it was added.
+
+  // The child graph depends on the events to get the frozen data
+  ::std::vector<cudaGraphNode_t> fX_ready_nodes = reserved::join_with_graph_nodes(ctx, fX_get_events, ctx.stage());
+
+  cudaGraphNodeParams cParams = {};
+  cParams.type                = cudaGraphNodeTypeConditional;
+  cParams.conditional.handle  = handle;
+  cParams.conditional.type    = cudaGraphCondTypeWhile;
+  cParams.conditional.size    = 1;
+
+  cudaGraphNode_t conditionalNode;
+#  if _CCCL_CTK_AT_LEAST(13, 0)
+  cudaGraphAddNode(&conditionalNode, ctx.get_graph(), fX_ready_nodes.data(), nullptr, fX_ready_nodes.size(), &cParams);
+#  else
+  cudaGraphAddNode(&conditionalNode, ctx.get_graph(), fX_ready_nodes.data(), fX_ready_nodes.size(), &cParams);
+#  endif
+
+  cudaGraph_t bodyGraph = cParams.conditional.phGraph_out[0];
+
+  // A child graph contains the entire body
+  cudaGraphNode_t child_graph_node;
+  cuda_safe_call(cudaGraphAddChildGraphNode(&child_graph_node, bodyGraph, nullptr, 0, sub_ctx.get_graph()));
+
+  // Create an event that depends on the conditional node, so that we unfreeze
+  // after the completion of the while loop
+  event_list child_graph_event;
+  reserved::fork_from_graph_node(
+    ctx, conditionalNode, ctx.get_graph(), ctx.stage(), child_graph_event, "child graph done");
+  fX.unfreeze(child_graph_event);
+
+  ctx.host_launch(lX.read())->*[](auto x) {
+    for (int i = 0; i < static_cast<int>(x.size()); i++)
+    {
+      EXPECT(x(i) == 3 * X0(i) + 2 * 5);
+    }
+  };
+
+  ctx.finalize();
+#endif // !_CCCL_CTK_BELOW(12, 4)
+}

cccl_upstream/cudax/test/stf/graph/freeze_while_graphs_2.cu

@@ -0,0 +1,128 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2025 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+//! \file
+//!
+//! \brief Freeze a logical data in a graph to use it in the body of a "while" graph node, the resulting looping graph
+//! will be executed within a stream context.
+
+#include <cuda/experimental/stf.cuh>
+
+#include <vector>
+
+using namespace cuda::experimental::stf;
+
+#if _CCCL_CTK_AT_LEAST(12, 4)
+int X0(int i)
+{
+  return 17 * i + 45;
+}
+
+__global__ void setHandle(cudaGraphConditionalHandle handle)
+{
+  static int count = 5;
+  cudaGraphSetConditional(handle, --count ? 1 : 0);
+}
+
+#endif // _CCCL_CTK_AT_LEAST(12, 4)
+
+int main()
+{
+#if _CCCL_CTK_BELOW(12, 4)
+  fprintf(stderr, "Skipping test: conditional nodes are only available since CUDA 12.4.\n");
+#else
+  const int N = 16;
+  int X[N];
+
+  for (int i = 0; i < N; i++)
+  {
+    X[i] = X0(i);
+  }
+
+  stream_ctx ctx;
+
+  auto lX = ctx.logical_data(X);
+
+  ctx.parallel_for(lX.shape(), lX.rw())->*[] __device__(size_t i, auto x) {
+    x(i) *= 3;
+  };
+
+  /* We are going to create a local context which is a graph, and we will populate it using a graph_ctx */
+  cudaGraph_t graph;
+  cuda_safe_call(cudaGraphCreate(&graph, 0));
+
+  cudaGraphConditionalHandle handle;
+  cudaGraphConditionalHandleCreate(&handle, graph, 1, cudaGraphCondAssignDefault);
+
+  // Create a graph that will later be inserted as a child graph once all input
+  // dependencies are known
+  cudaGraph_t sub_graph;
+  cuda_safe_call(cudaGraphCreate(&sub_graph, 0));
+
+  cudaGraphNodeParams cParams = {};
+  cParams.type                = cudaGraphNodeTypeConditional;
+  cParams.conditional.handle  = handle;
+  cParams.conditional.type    = cudaGraphCondTypeWhile;
+  cParams.conditional.size    = 1;
+
+  cudaGraphNode_t conditionalNode;
+  // There is no input dependency because they are implied by graph launch
+#  if _CCCL_CTK_AT_LEAST(13, 0)
+  cuda_safe_call(cudaGraphAddNode(&conditionalNode, graph, nullptr, nullptr, 0, &cParams));
+#  else
+  cuda_safe_call(cudaGraphAddNode(&conditionalNode, graph, nullptr, 0, &cParams));
+#  endif
+
+  cudaGraph_t bodyGraph = cParams.conditional.phGraph_out[0];
+
+  // Create a context based on this child graph which is the body of the
+  graph_ctx sub_ctx(bodyGraph);
+
+  auto fX                        = ctx.freeze(lX, access_mode::rw, data_place::current_device());
+  auto [frozen_X, fX_get_events] = fX.get(data_place::current_device());
+
+  auto lX_alias = sub_ctx.logical_data(frozen_X, data_place::current_device());
+
+  sub_ctx.parallel_for(lX.shape(), lX_alias.rw())->*[] __device__(size_t i, auto x) {
+    x(i) = x(i) + 2;
+  };
+
+  // We want to repeat this a fixed number of times
+  sub_ctx.cuda_kernel()->*[handle]() {
+    return cuda_kernel_desc{setHandle, 1, 1, 0, handle};
+  };
+
+  sub_ctx.finalize_as_graph();
+
+  // The sub graph depends on the events to get the frozen data, so the
+  // launch of the graph will depend on them
+
+  cudaGraphExec_t graph_exec = NULL;
+  cuda_safe_call(cudaGraphInstantiate(&graph_exec, graph, NULL, NULL, 0));
+
+  auto support_dstream = ctx.pick_dstream();
+  /* auto before_launch = */ reserved::join_with_stream(ctx, support_dstream, fX_get_events, "graph_launch", false);
+
+  cuda_safe_call(cudaGraphLaunch(graph_exec, support_dstream.stream));
+
+  // We stop using the frozen logical data after then graph has been launched
+  auto graph_launched = reserved::record_event_in_stream(support_dstream);
+  fX.unfreeze(event_list(mv(graph_launched)));
+
+  ctx.host_launch(lX.read())->*[](auto x) {
+    for (int i = 0; i < static_cast<int>(x.size()); i++)
+    {
+      EXPECT(x(i) == 3 * X0(i) + 2 * 5);
+    }
+  };
+
+  ctx.finalize();
+#endif // !_CCCL_CTK_BELOW(12, 4)
+}

cccl_upstream/cudax/test/stf/graph/freeze_while_graphs_3.cu

@@ -0,0 +1,146 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2025 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+//! \file
+//!
+//! \brief Freeze a logical data in a graph to use it in the body of a "while" graph node, the resulting looping graph
+//! will be executed within a stream context.
+
+#include <cuda/experimental/stf.cuh>
+
+#include <vector>
+
+using namespace cuda::experimental::stf;
+
+#if _CCCL_CTK_AT_LEAST(12, 4)
+
+/**
+ * @brief Insert an existing CUDA graph node into a graph context with appropriate dependencies
+ *
+ * This function is designed for graph contexts and adds the provided graph node
+ * to the context's graph with dependencies from the input prerequisites.
+ *
+ * @tparam ctx_t Context type (must be a graph_ctx or context using graph_ctx under the hood)
+ * @param ctx The execution context (must be a graph context)
+ * @param node The existing CUDA graph node to insert
+ * @param input_prereqs Input dependencies that must be satisfied (must be graph events)
+ * @return event_list Events representing the completion of the graph node insertion
+ */
+template <typename ctx_t>
+event_list insert_graph_node(ctx_t& ctx, cudaGraphNode_t node, event_list& input_prereqs)
+{
+  cudaGraph_t support_graph = ctx.graph();
+  size_t graph_stage        = ctx.stage();
+
+  ::std::vector<cudaGraphNode_t> ready_nodes = reserved::join_with_graph_nodes(ctx, input_prereqs, graph_stage);
+
+  // Add dependencies from the ready_nodes to the existing node
+  if (!ready_nodes.empty())
+  {
+#  if _CCCL_CTK_AT_LEAST(13, 0)
+    cuda_safe_call(cudaGraphAddDependencies(support_graph, ready_nodes.data(), &node, nullptr, ready_nodes.size()));
+#  else // _CCCL_CTK_AT_LEAST(13, 0)
+    cuda_safe_call(cudaGraphAddDependencies(support_graph, ready_nodes.data(), &node, ready_nodes.size()));
+#  endif // _CCCL_CTK_AT_LEAST(13, 0)
+  }
+
+  // Create an event that depends on the inserted graph node
+  auto node_event = reserved::graph_event(node, graph_stage, support_graph);
+  node_event->set_symbol(ctx, "inserted_graph_node");
+
+  // Return the event list from that single event
+  return event_list(mv(node_event));
+}
+
+int X0(int i)
+{
+  return 17 * i + 45;
+}
+
+__global__ void setHandle(cudaGraphConditionalHandle handle)
+{
+  static int count = 5;
+  cudaGraphSetConditional(handle, --count ? 1 : 0);
+}
+
+#endif // _CCCL_CTK_AT_LEAST(12, 4)
+
+int main()
+{
+#if _CCCL_CTK_BELOW(12, 4)
+  fprintf(stderr, "Waiving test: conditional nodes are only available since CUDA 12.4.\n");
+#else
+  const int N = 16;
+  int X[N];
+
+  for (int i = 0; i < N; i++)
+  {
+    X[i] = X0(i);
+  }
+
+  graph_ctx ctx;
+
+  auto lX = ctx.logical_data(X);
+
+  ctx.parallel_for(lX.shape(), lX.rw())->*[] __device__(size_t i, auto x) {
+    x(i) *= 3;
+  };
+
+  cudaGraphConditionalHandle handle;
+  cuda_safe_call(cudaGraphConditionalHandleCreate(&handle, ctx.graph(), 1, cudaGraphCondAssignDefault));
+
+  cudaGraphNodeParams cParams = {};
+  cParams.type                = cudaGraphNodeTypeConditional;
+  cParams.conditional.handle  = handle;
+  cParams.conditional.type    = cudaGraphCondTypeWhile;
+  cParams.conditional.size    = 1;
+
+  cudaGraphNode_t conditionalNode;
+  // There is no input dependencies yet, we will add them later
+#  if _CCCL_CTK_AT_LEAST(13, 0)
+  cuda_safe_call(cudaGraphAddNode(&conditionalNode, ctx.graph(), nullptr, nullptr, 0, &cParams));
+#  else
+  cuda_safe_call(cudaGraphAddNode(&conditionalNode, ctx.graph(), nullptr, 0, &cParams));
+#  endif
+
+  cudaGraph_t bodyGraph = cParams.conditional.phGraph_out[0];
+
+  graph_ctx sub_ctx(bodyGraph);
+
+  auto fX                        = ctx.freeze(lX, access_mode::rw, data_place::current_device());
+  auto [frozen_X, fX_get_events] = fX.get(data_place::current_device());
+
+  auto lX_alias = sub_ctx.logical_data(frozen_X, data_place::current_device());
+
+  sub_ctx.parallel_for(lX.shape(), lX_alias.rw())->*[] __device__(size_t i, auto x) {
+    x(i) = x(i) + 2;
+  };
+
+  // We want to repeat this a fixed number of times
+  sub_ctx.cuda_kernel()->*[handle]() {
+    return cuda_kernel_desc{setHandle, 1, 1, 0, handle};
+  };
+
+  sub_ctx.finalize_as_graph();
+
+  event_list cond_graph_launched = insert_graph_node(ctx, conditionalNode, fX_get_events);
+
+  fX.unfreeze(mv(cond_graph_launched));
+
+  ctx.host_launch(lX.read())->*[](auto x) {
+    for (int i = 0; i < static_cast<int>(x.size()); i++)
+    {
+      EXPECT(x(i) == 3 * X0(i) + 2 * 5);
+    }
+  };
+
+  ctx.finalize();
+#endif // !_CCCL_CTK_BELOW(12, 4)
+}

cccl_upstream/cudax/test/stf/graph/get_cache_stats.cu

@@ -0,0 +1,51 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief An example to query statistics about graph instantiation
+ */
+
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+int main()
+{
+  async_resources_handle handle;
+  for (size_t i = 0; i < 10; i++)
+  {
+    graph_ctx ctx(handle);
+    auto lA = ctx.logical_data(shape_of<slice<size_t>>(64));
+    ctx.launch(lA.write())->*[] _CCCL_DEVICE(auto t, slice<size_t> A) {
+      for (auto i : t.apply_partition(shape(A)))
+      {
+        A(i) = 2 * i;
+      }
+    };
+    ctx.finalize();
+
+    // Query statistics about the graph context : the first iteration needs to
+    // instantiate the graph, then we will reuse graphs saved in the handle.
+    auto* st = ctx.graph_get_cache_stat();
+    if (i == 0)
+    {
+      EXPECT(st->instantiate_cnt == 1);
+      EXPECT(st->update_cnt == 0);
+    }
+    else
+    {
+      EXPECT(st->instantiate_cnt == 0);
+      EXPECT(st->update_cnt == 1);
+    }
+
+    // fprintf(stderr, "nnodes %ld nedges %ld\n", st->nnodes, st->nedges);
+  }
+}

cccl_upstream/cudax/test/stf/graph/graph_cache_policy.cu

@@ -0,0 +1,57 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief An example to query statistics about graph instantiation
+ */
+
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+int main()
+{
+  async_resources_handle handle;
+  for (size_t i = 0; i < 10; i++)
+  {
+    graph_ctx ctx(handle);
+
+    // If i is a multiple of 3 we enable the cache, the first iteration will fill the cache
+    ctx.set_graph_cache_policy([i]() {
+      return (i % 3) == 0;
+    });
+
+    auto lA = ctx.logical_data(shape_of<slice<size_t>>(64));
+    ctx.launch(lA.write())->*[] _CCCL_DEVICE(auto t, slice<size_t> A) {
+      for (auto i : t.apply_partition(shape(A)))
+      {
+        A(i) = 2 * i;
+      }
+    };
+    ctx.finalize();
+
+    // Query statistics about the graph context : the first iteration needs to
+    // instantiate the graph, then we will reuse graphs saved in the handle.
+    auto* st = ctx.graph_get_cache_stat();
+
+    // For the first iteration, or non multiple of 3 we have to instantiate, otherwise we should have a cache hit
+    if (i == 0 || (i % 3) != 0)
+    {
+      EXPECT(st->instantiate_cnt == 1);
+      EXPECT(st->update_cnt == 0);
+    }
+    else
+    {
+      EXPECT(st->instantiate_cnt == 0);
+      EXPECT(st->update_cnt == 1);
+    }
+  }
+}

cccl_upstream/cudax/test/stf/graph/graph_child_move_ownership.cu

@@ -0,0 +1,68 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2026 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief Verify that graph_ctx tasks whose captured child graphs contain
+ *        memory allocation/free nodes (from cudaMallocAsync) work correctly.
+ *
+ * Before the move-ownership fix, cudaGraphAddChildGraphNode (clone semantics)
+ * rejected such child graphs with CUDA_ERROR_NOT_SUPPORTED.  CTK 13+ exposes
+ * cudaGraphChildGraphOwnershipMove via cudaGraphAddNode which transfers the
+ * child graph to the parent instead of cloning it.
+ */
+
+#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+__global__ void fill_kernel(int* ptr, int n, int val)
+{
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  if (tid < n)
+  {
+    ptr[tid] = val;
+  }
+}
+
+int main()
+{
+#if _CCCL_CTK_BELOW(13, 0)
+  fprintf(stderr, "Waiving test: cudaGraphChildGraphOwnershipMove requires CTK 13+.\n");
+#else
+  constexpr int N = 256;
+  int host_data[N];
+  for (int i = 0; i < N; i++)
+  {
+    host_data[i] = 0;
+  }
+
+  graph_ctx ctx;
+  auto ldata = ctx.logical_data(host_data);
+
+  // The lambda receives cudaStream_t, so graph_ctx uses stream capture.
+  // cudaMallocAsync/cudaFreeAsync on that stream produce mem-alloc/free
+  // graph nodes inside the captured child graph.
+  ctx.task(ldata.rw())->*[](cudaStream_t s, auto sdata) {
+    int* tmp = nullptr;
+    cuda_safe_call(cudaMallocAsync(&tmp, N * sizeof(int), s));
+    fill_kernel<<<(N + 255) / 256, 256, 0, s>>>(tmp, N, 42);
+    cuda_safe_call(cudaMemcpyAsync(sdata.data_handle(), tmp, N * sizeof(int), cudaMemcpyDeviceToDevice, s));
+    cuda_safe_call(cudaFreeAsync(tmp, s));
+  };
+
+  ctx.finalize();
+
+  for (int i = 0; i < N; i++)
+  {
+    assert(host_data[i] == 42);
+  }
+#endif // !_CCCL_CTK_BELOW(13, 0)
+}

cccl_upstream/cudax/test/stf/graph/graph_composition.cu

@@ -0,0 +1,117 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief Generate a library call from nested CUDA graphs generated using algorithms
+ */
+
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+// Some fake library doing MATH
+void libMATH(graph_ctx ctx, logical_data<slice<double>> x, logical_data<slice<double>> y)
+{
+  // We only want to have kernels with 4 CTAs to stress the system
+  auto spec = par<4>(par<128>());
+  ctx.launch(spec, exec_place::current_device(), x.read(), y.write()).set_symbol("MATH1")->*
+    [] __device__(auto t, auto x, auto y) {
+      for (auto i : t.apply_partition(shape(x)))
+      {
+        y(i) = cos(cos(x(i)));
+      }
+    };
+
+  ctx.launch(spec, exec_place::current_device(), x.write(), y.read()).set_symbol("MATH2")->*
+    [] __device__(auto t, auto x, auto y) {
+      for (auto i : t.apply_partition(shape(x)))
+      {
+        x(i) = sin(sin(y(i)));
+      };
+    };
+}
+
+template <typename context_t>
+void libMATH_AS_GRAPH(context_t& ctx, logical_data<slice<double>> x, logical_data<slice<double>> y)
+{
+  static algorithm alg;
+  alg.run_as_task(libMATH, ctx, x.rw(), y.write());
+}
+
+// Some fake lib doing a SWAP
+template <typename context_t>
+void libSWAP(context_t& ctx, logical_data<slice<double>> x, logical_data<slice<double>> y)
+{
+  // We only want to have kernels with 4 CTAs to stress the system
+  auto spec = par<4>(par<128>());
+  ctx.launch(spec, exec_place::current_device(), x.rw(), y.rw()).set_symbol("SWAP")->*
+    [] __device__(auto t, auto x, auto y) {
+      for (auto i : t.apply_partition(shape(x)))
+      {
+        auto tmp = x(i);
+        x(i)     = y(i);
+        y(i)     = tmp;
+      }
+    };
+}
+
+template <typename context_t>
+logical_data<slice<double>> libCOPY(context_t& ctx, logical_data<slice<double>> x)
+{
+  logical_data<slice<double>> res = ctx.logical_data(x.shape());
+
+  // We only want to have kernels with 4 CTAs to stress the system
+  auto spec = par<4>(par<128>());
+  ctx.launch(spec, exec_place::current_device(), x.read(), res.write()).set_symbol("SWAP")->*
+    [] __device__(auto t, auto x, auto res) {
+      for (auto i : t.apply_partition(shape(x)))
+      {
+        res(i) = x(i);
+      }
+    };
+
+  return res;
+}
+
+int main()
+{
+  nvtx_range r("run");
+
+  stream_ctx ctx;
+
+  const size_t N = 256 * 1024;
+  const size_t K = 8;
+
+  logical_data<slice<double>> lX[K];
+  logical_data<slice<double>> lY[K];
+
+  for (size_t i = 0; i < K; i++)
+  {
+    lX[i] = ctx.logical_data<double>(N);
+    lY[i] = ctx.logical_data<double>(N);
+
+    ctx.parallel_for(lX[i].shape(), lX[i].write(), lY[i].write()).set_symbol("INIT")->*
+      [] __device__(size_t i, auto x, auto y) {
+        x(i) = 2.0 * i + 12.0;
+        y(i) = -3.0 * i + 17.0;
+      };
+  }
+
+  for (size_t i = 0; i < K; i++)
+  {
+    auto tmp = libCOPY(ctx, lX[i]);
+    libSWAP(ctx, tmp, lY[i]);
+    libMATH_AS_GRAPH(ctx, lX[i], lY[i]);
+    libSWAP(ctx, lX[i], lY[i]);
+  }
+
+  ctx.finalize();
+}