cuda/src/tbb_await.cpp

0001 /** TRACCC library, part of the ACTS project (R&D line)
0002  *
0003  * (c) 2026 CERN for the benefit of the ACTS project
0004  *
0005  * Mozilla Public License Version 2.0
0006  */
0007
0008 // Local include(s).
0009 #include "traccc/examples/cuda/tbb_await.hpp"
0010
0011 #include "cuda_error_check.hpp"
0012
0013 // CUDA include(s).
0014 #include <cuda_runtime_api.h>
0015
0016 // TBB include(s).
0017 #include <tbb/task.h>
0018
0019 namespace traccc::cuda {
0020
0021 namespace {
0022 void CUDART_CB suspend_stream_callback(void* tag) {
0023   tbb::task::resume(*static_cast<tbb::task::suspend_point*>(tag));
0024 }
0025 }  // namespace
0026
0027 void tbb_await_callback(vecmem::abstract_event& event,
0028                         const stream_wrapper& stream) {
0029   event.ignore();  // ignore the event, as it is not needed for resumption
0030   cudaError_t err = cudaSuccess;
0031   auto suspend_point =
0032       tbb::task::suspend_point{};  // suspension point address must remain valid
0033                                    // when resumption callback is called
0034   tbb::task::suspend([&err, &stream, &suspend_point](auto tag) {
0035     suspend_point = tag;
0036     auto cuda_stream = reinterpret_cast<cudaStream_t>(stream.cudaStream());
0037     err = cudaLaunchHostFunc(cuda_stream, suspend_stream_callback,
0038                              &suspend_point);
0039     // resume immediately if the callback could not be registered
0040     if (err != cudaSuccess) {
0041       tbb::task::resume(suspend_point);
0042     }
0043   });
0044   CUDA_ERROR_CHECK(err);
0045 }
0046 }  // namespace traccc::cuda