add "reco" examples

CMakeLists.txt

@@ -2,6 +2,7 @@ cmake_minimum_required(VERSION 3.25)
 project(CoroutineTests)
 
 set(CMAKE_CXX_STANDARD 20 CACHE STRING "")
+set(CMAKE_CXX_EXTENSIONS FALSE CACHE BOOL "Disable C++ extensions" )
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 if(CMAKE_CXX_STANDARD LESS 20)
   message(FATAL_ERROR "Unsupported C++ standard: ${CMAKE_CXX_STANDARD}. Must be at least C++20.")
@@ -81,6 +82,13 @@ if(BUILD_TBB)
   endif()
 endif()
 
+### CUDA options
+
+option(BUILD_CUDA "Build with CUDA support" OFF)
+if(BUILD_CUDA)
+  find_package(CUDAToolkit REQUIRED)
+endif()
+
 ## targets
 
 add_library(CoroutineTests INTERFACE)

CMakePresets.json

@@ -12,7 +12,8 @@
                 "BUILD_STDEXEC": "ON",
                 "BUILD_TBB": "ON",
                 "BUILD_EXAMPLES": "ON",
-                "BUILD_PERF": "ON"
+                "BUILD_PERF": "ON",
+                "BUILD_CUDA": "OFF"
             }
         },
         {
@@ -26,6 +27,16 @@
                 "BUILD_BEMAN_TASK": "OFF"
             }
         },
+        {
+            "name": "20-cuda",
+            "displayName": "Configure preset constrained to C++20 with CUDA support",
+            "inherits": [
+                "20"
+            ],
+            "cacheVariables": {
+                "BUILD_CUDA": "ON"
+            }
+        },
         {
             "name": "23",
             "displayName": "Configure preset constrained to C++23",
@@ -39,7 +50,8 @@
         }
     ],
     "buildPresets": [
-        {   "name": "default",
+        {
+            "name": "default",
             "displayName": "Default Build Preset",
             "hidden": true,
             "configurePreset": "default",
@@ -53,6 +65,14 @@
                 "default"
             ]
         },
+        {
+            "name": "20-cuda",
+            "displayName": "Build with C++20 and CUDA support",
+            "configurePreset": "20-cuda",
+            "inherits": [
+                "20"
+            ]
+        },
         {
             "name": "23",
             "displayName": "Build with C++23",
@@ -76,6 +96,19 @@
                 }
             ]
         },
+        {
+            "name": "20-cuda",
+            "steps": [
+                {
+                    "type": "configure",
+                    "name": "20-cuda"
+                },
+                {
+                    "type": "build",
+                    "name": "20-cuda"
+                }
+            ]
+        },
         {
             "name": "23",
             "steps": [

examples/CMakeLists.txt

@@ -55,3 +55,12 @@ if(BUILD_BEMAN_TASK)
     target_compile_definitions(exec_tbb_beman PUBLIC "USE_BEMAN")
   endif()
 endif()
+
+if(BUILD_CUDA AND BUILD_TBB)
+  add_executable(alien_reco alien_reco.cpp)
+  target_link_libraries(alien_reco PRIVATE CoroutineTests CUDA::cudart TBB::tbb)
+  if(BUILD_STDEXEC)
+    add_executable(exec_reco_stdexec exec_reco.cpp)
+    target_link_libraries(exec_reco_stdexec PRIVATE CoroutineTests CUDA::cudart TBB::tbb STDEXEC::stdexec)
+  endif()
+endif()

examples/README.md

@@ -143,3 +143,9 @@ This example demonstrates a `sync_wait` algorithm compatible with coroutine sema
 Link: [alien_counting_scope.cpp](alien_counting_scope.cpp)
 
 This example demonstrates dynamic work submitting with `counting_scope` compatible with coroutine semantic as in ["Alien" example](#alien). `spawn` schedules execution of a coroutine that returns `void`. `join()` blocks the current thread until all submitted coroutines are finished (and rethrows the first exception captured from submitted work, if any).
+
+## Reconstruction
+
+Link: [alien_reco.cpp](alien_reco.cpp) and [exec_reco.cpp](exec_reco.cpp)
+
+These examples show a setup and coroutine chain loosely inspired by track reconstruction on GPU in high energy physics experiments. The `reconstruct` coroutine prepares a mockup input data on a CUDA device, then `co_await`s the `clustering` coroutine, then receives output from it and `co_await`s the `seeding` coroutine. Both `clustering` and `seeding` coroutines receive a device buffer, copy it asynchronously back to host, suspend until the copy is done, then count non-zero elements and allocate new buffer of that size for their results. In `main` the `reconstruct` coroutines are executed in a TBB task arena either synchronously waiting for the result from the submitting or dynamically starting a few `reconstruct` coroutines and waiting until all the work is finished. In the `alien_reco` the application is implemented with "alien" coroutines (as in [alien examples](#alien)) , while in `exec_reco_stdexec` the application is implemented with C++26 execution.

examples/alien_reco.cpp

@@ -0,0 +1,237 @@
+#include <cuda_runtime_api.h>
+#include <tbb/task_arena.h>
+
+#include <cstddef>
+#include <iostream>
+
+#include "CoroutineTests/alien/counting_scope.hpp"
+#include "CoroutineTests/alien/subtool.hpp"
+#include "CoroutineTests/alien/sync_wait.hpp"
+#include "CoroutineTests/alien/tool.hpp"
+#include "logging_utils.hpp"  // log, format_name
+
+#define ERROR_CHECK_CUDA(EXP)                                              \
+    do {                                                                   \
+        cudaError_t errorCode = EXP;                                       \
+        if (errorCode != cudaSuccess) {                                    \
+            throw std::runtime_error(                                      \
+                std::format("CUDA error at {}:{}: {}", __FILE__, __LINE__, \
+                            cudaGetErrorString(errorCode)));               \
+        }                                                                  \
+    } while (false)
+
+using namespace CoroutineTests::alien;
+
+/// Awaitable that resumes a coroutine when a CUDA stream reaches a certain
+/// point. Internally cudaLaunchHostFunc is used to set up a resumption callback
+/// on the stream.
+class StreamAwaitable {
+    public:
+    StreamAwaitable(cudaStream_t stream) : m_stream(stream) {}
+
+    bool await_ready() const noexcept { return false; }
+    template <typename Promise>
+    void await_suspend(std::coroutine_handle<Promise> handle) {
+        m_error = cudaLaunchHostFunc(m_stream, resumption_callback<Promise>,
+                                     handle.address());
+        // If the callback couldn't be registered, we need to reschedule the
+        // coroutine immediately to avoid deadlock.
+        if (m_error != cudaSuccess) {
+            handle.promise().reschedule();
+        }
+    }
+    cudaError_t await_resume() const noexcept { return m_error; }
+
+    private:
+    cudaStream_t m_stream;
+    cudaError_t m_error = cudaSuccess;
+
+    template <typename Promise>
+    static void resumption_callback(void* userData) {
+        auto handle = std::coroutine_handle<Promise>::from_address(userData);
+        handle.promise().reschedule();
+    }
+};
+
+template <typename T>
+struct DeviceBuffer {
+    T* ptr = nullptr;
+    std::size_t size = 0;
+};
+
+subtool::Task<DeviceBuffer<int>> clusterization(DeviceBuffer<int> cells,
+                                                cudaStream_t stream,
+                                                std::string_view parent) {
+
+    const auto self = format_name(parent, "clusterization");
+    log(self) << "Starting clusterization" << std::endl;
+
+    const auto nCells = static_cast<int>(cells.size);
+
+    // Copy cells back to host to count non-zero entries
+    auto h_cells = std::vector<int>(nCells);
+
+    ERROR_CHECK_CUDA(cudaMemcpyAsync(h_cells.data(), cells.ptr,
+                                     nCells * sizeof(int),
+                                     cudaMemcpyDeviceToHost, stream));
+
+    ERROR_CHECK_CUDA(co_await StreamAwaitable{stream});
+
+    auto nClusters = 0;
+    for (auto v : h_cells)
+        if (v != 0)
+            ++nClusters;
+
+    log(self) << "Found " << nClusters << " clusters" << std::endl;
+
+    // Allocate clusters of appropiate size on device
+    int* d_clusters = nullptr;
+    ERROR_CHECK_CUDA(cudaMallocAsync(reinterpret_cast<void**>(&d_clusters),
+                                     nClusters * sizeof(int), stream));
+
+    // Write some dummy data to the clusters buffer to simulate work
+    ERROR_CHECK_CUDA(
+        cudaMemsetAsync(d_clusters, 0, nClusters * sizeof(int), stream));
+    ERROR_CHECK_CUDA(
+        cudaMemsetAsync(d_clusters, 1, nClusters / 2 * sizeof(int), stream));
+
+    co_return DeviceBuffer<int>{d_clusters,
+                                static_cast<std::size_t>(nClusters)};
+}
+
+subtool::Task<DeviceBuffer<int>> seeding(DeviceBuffer<int> clusters,
+                                         cudaStream_t stream,
+                                         std::string_view parent) {
+
+    const auto self = format_name(parent, "seeding");
+    log(self) << "Starting seeding" << std::endl;
+
+    const auto nClusters = static_cast<int>(clusters.size);
+
+    // Copy clusters to host to count non-zero entries
+    auto h_clusters = std::vector<int>(nClusters);
+
+    ERROR_CHECK_CUDA(cudaMemcpyAsync(h_clusters.data(), clusters.ptr,
+                                     nClusters * sizeof(int),
+                                     cudaMemcpyDeviceToHost, stream));
+
+    ERROR_CHECK_CUDA(co_await StreamAwaitable{stream});
+
+    int nSeeds = 0;
+    for (auto v : h_clusters)
+        if (v != 0)
+            ++nSeeds;
+
+    log(self) << "Found " << nSeeds << " seeds" << std::endl;
+
+    // Allocate clusters of appropiate size on device
+    int* d_seeds = nullptr;
+    ERROR_CHECK_CUDA(cudaMallocAsync(reinterpret_cast<void**>(&d_seeds),
+                                     nSeeds * sizeof(int), stream));
+
+    // Write some dummy data to the clusters buffer to simulate work
+    ERROR_CHECK_CUDA(cudaMemsetAsync(d_seeds, 0, nSeeds * sizeof(int), stream));
+    ERROR_CHECK_CUDA(
+        cudaMemsetAsync(d_seeds, 1, nSeeds / 2 * sizeof(int), stream));
+
+    co_return DeviceBuffer<int>{d_seeds, static_cast<std::size_t>(nSeeds)};
+}
+
+tool::Task<tool::StatusCode> reconstruct(cudaStream_t stream,
+                                         std::string_view parent) {
+    const auto self = format_name(parent, "reconstruction");
+    log(self) << "Starting reconstruction" << std::endl;
+
+    // Allocate some dummy input data on the device
+    auto cells = DeviceBuffer<int>{nullptr, 1000};
+    ERROR_CHECK_CUDA(cudaMallocAsync(reinterpret_cast<void**>(&cells.ptr),
+                                     cells.size * sizeof(int), stream));
+    ERROR_CHECK_CUDA(
+        cudaMemsetAsync(cells.ptr, 1, cells.size * sizeof(int), stream));
+
+    // Run the clusterization and seeding steps
+    auto clusters = co_await clusterization(cells, stream, self);
+    auto seeds = co_await seeding(clusters, stream, self);
+
+    // Cleanup
+    ERROR_CHECK_CUDA(cudaFreeAsync(cells.ptr, stream));
+    ERROR_CHECK_CUDA(cudaFreeAsync(clusters.ptr, stream));
+    ERROR_CHECK_CUDA(cudaFreeAsync(seeds.ptr, stream));
+
+    ERROR_CHECK_CUDA(co_await StreamAwaitable{stream});
+
+    log(self) << "Finishing reconstruction" << std::endl;
+    co_return tool::StatusCode::SUCCESS;
+}
+
+int main() {
+    int deviceCount = 0;
+    auto error_id = cudaGetDeviceCount(&deviceCount);
+
+    if (error_id != cudaSuccess) {
+        std::cout << "cudaGetDeviceCount returned "
+                  << static_cast<int>(error_id) << "\n"
+                  << cudaGetErrorString(error_id) << "\n";
+        return EXIT_FAILURE;
+    }
+
+    if (deviceCount == 0) {
+        std::cout << "No CUDA devices found.\n";
+        return EXIT_FAILURE;
+    }
+
+    log() << "main Starting" << std::endl;
+
+    tbb::task_arena task_arena{2};
+
+    auto scheduler = [&task_arena](std::coroutine_handle<> handle) {
+        task_arena.enqueue([handle]() { handle.resume(); });
+    };
+
+    {
+        cudaStream_t stream;
+        ERROR_CHECK_CUDA(cudaStreamCreate(&stream));
+        std::cout << "--- Single event, synchronous wait for completion ---\n";
+        log() << "main Launching algorithm..." << std::endl;
+        auto status = sync_wait(scheduler, reconstruct(stream, "main"));
+        log() << "main Final status of algorithm " << status << "" << std::endl;
+        ERROR_CHECK_CUDA(cudaStreamDestroy(stream));
+    }
+
+    {
+        std::cout << "--- Multiple events, wait for all to complete ---\n";
+
+        auto streams = std::vector<cudaStream_t>(2);
+        auto status = std::vector<tool::StatusCode>(streams.size());
+        for (auto& stream : streams) {
+            ERROR_CHECK_CUDA(cudaStreamCreate(&stream));
+        }
+
+        auto scope = counting_scope{};
+        log() << "main Launching algorithms..." << std::endl;
+
+        auto payload = [](std::vector<cudaStream_t> streams,
+                          std::vector<tool::StatusCode>& statuses,
+                          int i) -> tool::Task<void> {
+            const auto name = std::format("event{}:main", i);
+            auto& stream = streams.at(i);
+            auto& status = statuses.at(i);
+            status = co_await reconstruct(stream, name);
+        };
+        for (std::size_t i = 0; i < streams.size(); ++i) {
+            scope.spawn(scheduler, payload(streams, status, i));
+        }
+        scope.join();
+
+        for (auto& stream : streams) {
+            ERROR_CHECK_CUDA(cudaStreamDestroy(stream));
+        }
+
+        log() << "main All algorithms completed. Final statuses: ";
+        for (auto s : status) {
+            std::cout << s << ' ';
+        }
+        std::cout << std::endl;
+    }
+    return 0;
+}