hip/src/tbb_await.cpp

0001 // This file is part of the ACTS project.
0002 //
0003 // Copyright (C) 2016 CERN for the benefit of the ACTS project
0004 //
0005 // This Source Code Form is subject to the terms of the Mozilla Public
0006 // License, v. 2.0. If a copy of the MPL was not distributed with this
0007 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
0008
0009 // Local include(s).
0010 #include "traccc/examples/hip/tbb_await.hpp"
0011
0012 #include "hip_error_check.hpp"
0013
0014 // HIP include(s).
0015 #include <hip/hip_runtime_api.h>
0016
0017 // TBB include(s).
0018 #include <tbb/task.h>
0019
0020 namespace traccc::hip {
0021
0022 namespace {
0023 void suspend_stream_callback(void* tag) {
0024   tbb::task::resume(*static_cast<tbb::task::suspend_point*>(tag));
0025 }
0026 }  // namespace
0027
0028 void tbb_await_callback(vecmem::abstract_event& event,
0029                         const stream_wrapper& stream) {
0030   event.ignore();  // ignore the event, as it is not needed for resumption
0031   hipError_t err = hipSuccess;
0032   auto suspend_point =
0033       tbb::task::suspend_point{};  // suspension point address must remain valid
0034                                    // when resumption callback is called
0035   tbb::task::suspend([&err, &stream, &suspend_point](auto tag) {
0036     suspend_point = tag;
0037     auto hip_stream = reinterpret_cast<hipStream_t>(stream.hipStream());
0038     err =
0039         hipLaunchHostFunc(hip_stream, suspend_stream_callback, &suspend_point);
0040     // resume immediately if the callback could not be registered
0041     if (err != hipSuccess) {
0042       tbb::task::resume(suspend_point);
0043     }
0044   });
0045   HIP_ERROR_CHECK(err);
0046 }
0047 }  // namespace traccc::hip