|
| 1 | +#include <cuda_runtime_api.h> |
| 2 | +#include <tbb/task_arena.h> |
| 3 | + |
| 4 | +#include <cstddef> |
| 5 | +#include <iostream> |
| 6 | + |
| 7 | +#include "CoroutineTests/alien/counting_scope.hpp" |
| 8 | +#include "CoroutineTests/alien/subtool.hpp" |
| 9 | +#include "CoroutineTests/alien/sync_wait.hpp" |
| 10 | +#include "CoroutineTests/alien/tool.hpp" |
| 11 | +#include "logging_utils.hpp" // log, format_name |
| 12 | + |
| 13 | +#define ERROR_CHECK_CUDA(EXP) \ |
| 14 | + do { \ |
| 15 | + cudaError_t errorCode = EXP; \ |
| 16 | + if (errorCode != cudaSuccess) { \ |
| 17 | + throw std::runtime_error( \ |
| 18 | + std::format("CUDA error at {}:{}: {}", __FILE__, __LINE__, \ |
| 19 | + cudaGetErrorString(errorCode))); \ |
| 20 | + } \ |
| 21 | + } while (false) |
| 22 | + |
| 23 | +using namespace CoroutineTests::alien; |
| 24 | + |
| 25 | +/// Awaitable that resumes a coroutine when a CUDA stream reaches a certain |
| 26 | +/// point. Internally cudaLaunchHostFunc is used to set up a resumption callback |
| 27 | +/// on the stream. |
| 28 | +class StreamAwaitable { |
| 29 | + public: |
| 30 | + StreamAwaitable(cudaStream_t stream) : m_stream(stream) {} |
| 31 | + |
| 32 | + bool await_ready() const noexcept { return false; } |
| 33 | + template <typename Promise> |
| 34 | + void await_suspend(std::coroutine_handle<Promise> handle) { |
| 35 | + m_error = cudaLaunchHostFunc(m_stream, resumption_callback<Promise>, |
| 36 | + handle.address()); |
| 37 | + // If the callback couldn't be registered, we need to reschedule the |
| 38 | + // coroutine immediately to avoid deadlock. |
| 39 | + if (m_error != cudaSuccess) { |
| 40 | + handle.promise().reschedule(); |
| 41 | + } |
| 42 | + } |
| 43 | + cudaError_t await_resume() const noexcept { return m_error; } |
| 44 | + |
| 45 | + private: |
| 46 | + cudaStream_t m_stream; |
| 47 | + cudaError_t m_error = cudaSuccess; |
| 48 | + |
| 49 | + template <typename Promise> |
| 50 | + static void resumption_callback(void* userData) { |
| 51 | + auto handle = std::coroutine_handle<Promise>:: from_address(userData); |
| 52 | + handle.promise().reschedule(); |
| 53 | + } |
| 54 | +}; |
| 55 | + |
| 56 | +template <typename T> |
| 57 | +struct DeviceBuffer { |
| 58 | + T* ptr = nullptr; |
| 59 | + std::size_t size = 0; |
| 60 | +}; |
| 61 | + |
| 62 | +subtool::Task<DeviceBuffer<int>> clusterization(DeviceBuffer<int> cells, |
| 63 | + cudaStream_t stream, |
| 64 | + std::string_view parent) { |
| 65 | + |
| 66 | + const auto self = format_name(parent, "clusterization"); |
| 67 | + log(self) << "Starting clusterization" << std::endl; |
| 68 | + |
| 69 | + const auto nCells = static_cast<int>(cells.size); |
| 70 | + |
| 71 | + // Copy cells back to host to count non-zero entries |
| 72 | + auto h_cells = std::vector<int>(nCells); |
| 73 | + |
| 74 | + ERROR_CHECK_CUDA(cudaMemcpyAsync(h_cells.data(), cells.ptr, |
| 75 | + nCells * sizeof(int), |
| 76 | + cudaMemcpyDeviceToHost, stream)); |
| 77 | + |
| 78 | + ERROR_CHECK_CUDA(co_await StreamAwaitable{stream}); |
| 79 | + |
| 80 | + auto nClusters = 0; |
| 81 | + for (auto v : h_cells) |
| 82 | + if (v != 0) |
| 83 | + ++nClusters; |
| 84 | + |
| 85 | + log(self) << "Found " << nClusters << " clusters" << std::endl; |
| 86 | + |
| 87 | + // Allocate clusters of appropiate size on device |
| 88 | + int* d_clusters = nullptr; |
| 89 | + ERROR_CHECK_CUDA(cudaMallocAsync(reinterpret_cast<void**>(&d_clusters), |
| 90 | + nClusters * sizeof(int), stream)); |
| 91 | + |
| 92 | + // Write some dummy data to the clusters buffer to simulate work |
| 93 | + ERROR_CHECK_CUDA( |
| 94 | + cudaMemsetAsync(d_clusters, 0, nClusters * sizeof(int), stream)); |
| 95 | + ERROR_CHECK_CUDA( |
| 96 | + cudaMemsetAsync(d_clusters, 1, nClusters / 2 * sizeof(int), stream)); |
| 97 | + |
| 98 | + co_return DeviceBuffer<int>{d_clusters, |
| 99 | + static_cast<std::size_t>(nClusters)}; |
| 100 | +} |
| 101 | + |
| 102 | +subtool::Task<DeviceBuffer<int>> seeding(DeviceBuffer<int> clusters, |
| 103 | + cudaStream_t stream, |
| 104 | + std::string_view parent) { |
| 105 | + |
| 106 | + const auto self = format_name(parent, "seeding"); |
| 107 | + log(self) << "Starting seeding" << std::endl; |
| 108 | + |
| 109 | + const auto nClusters = static_cast<int>(clusters.size); |
| 110 | + |
| 111 | + // Copy clusters to host to count non-zero entries |
| 112 | + auto h_clusters = std::vector<int>(nClusters); |
| 113 | + |
| 114 | + ERROR_CHECK_CUDA(cudaMemcpyAsync(h_clusters.data(), clusters.ptr, |
| 115 | + nClusters * sizeof(int), |
| 116 | + cudaMemcpyDeviceToHost, stream)); |
| 117 | + |
| 118 | + ERROR_CHECK_CUDA(co_await StreamAwaitable{stream}); |
| 119 | + |
| 120 | + int nSeeds = 0; |
| 121 | + for (auto v : h_clusters) |
| 122 | + if (v != 0) |
| 123 | + ++nSeeds; |
| 124 | + |
| 125 | + log(self) << "Found " << nSeeds << " seeds" << std::endl; |
| 126 | + |
| 127 | + // Allocate clusters of appropiate size on device |
| 128 | + int* d_seeds = nullptr; |
| 129 | + ERROR_CHECK_CUDA(cudaMallocAsync(reinterpret_cast<void**>(&d_seeds), |
| 130 | + nSeeds * sizeof(int), stream)); |
| 131 | + |
| 132 | + // Write some dummy data to the clusters buffer to simulate work |
| 133 | + ERROR_CHECK_CUDA(cudaMemsetAsync(d_seeds, 0, nSeeds * sizeof(int), stream)); |
| 134 | + ERROR_CHECK_CUDA( |
| 135 | + cudaMemsetAsync(d_seeds, 1, nSeeds / 2 * sizeof(int), stream)); |
| 136 | + |
| 137 | + co_return DeviceBuffer<int>{d_seeds, static_cast<std::size_t>(nSeeds)}; |
| 138 | +} |
| 139 | + |
| 140 | +tool::Task<tool::StatusCode> reconstruct(cudaStream_t stream, |
| 141 | + std::string_view parent) { |
| 142 | + const auto self = format_name(parent, "reconstruction"); |
| 143 | + log(self) << "Starting reconstruction" << std::endl; |
| 144 | + |
| 145 | + // Allocate some dummy input data on the device |
| 146 | + auto cells = DeviceBuffer<int>{nullptr, 1000}; |
| 147 | + ERROR_CHECK_CUDA(cudaMallocAsync(reinterpret_cast<void**>(&cells.ptr), |
| 148 | + cells.size * sizeof(int), stream)); |
| 149 | + ERROR_CHECK_CUDA( |
| 150 | + cudaMemsetAsync(cells.ptr, 1, cells.size * sizeof(int), stream)); |
| 151 | + |
| 152 | + // Run the clusterization and seeding steps |
| 153 | + auto clusters = co_await clusterization(cells, stream, self); |
| 154 | + auto seeds = co_await seeding(clusters, stream, self); |
| 155 | + |
| 156 | + // Cleanup |
| 157 | + ERROR_CHECK_CUDA(cudaFreeAsync(cells.ptr, stream)); |
| 158 | + ERROR_CHECK_CUDA(cudaFreeAsync(clusters.ptr, stream)); |
| 159 | + ERROR_CHECK_CUDA(cudaFreeAsync(seeds.ptr, stream)); |
| 160 | + |
| 161 | + ERROR_CHECK_CUDA(co_await StreamAwaitable{stream}); |
| 162 | + |
| 163 | + log(self) << "Finishing reconstruction" << std::endl; |
| 164 | + co_return tool::StatusCode::SUCCESS; |
| 165 | +} |
| 166 | + |
| 167 | +int main() { |
| 168 | + int deviceCount = 0; |
| 169 | + auto error_id = cudaGetDeviceCount(&deviceCount); |
| 170 | + |
| 171 | + if (error_id != cudaSuccess) { |
| 172 | + std::cout << "cudaGetDeviceCount returned " |
| 173 | + << static_cast<int>(error_id) << "\n" |
| 174 | + << cudaGetErrorString(error_id) << "\n"; |
| 175 | + return EXIT_FAILURE; |
| 176 | + } |
| 177 | + |
| 178 | + if (deviceCount == 0) { |
| 179 | + std::cout << "No CUDA devices found.\n"; |
| 180 | + return EXIT_FAILURE; |
| 181 | + } |
| 182 | + |
| 183 | + log() << "main Starting" << std::endl; |
| 184 | + |
| 185 | + tbb::task_arena task_arena{2}; |
| 186 | + |
| 187 | + auto scheduler = [&task_arena](std::coroutine_handle<> handle) { |
| 188 | + task_arena.enqueue([handle]() { handle.resume(); }); |
| 189 | + }; |
| 190 | + |
| 191 | + { |
| 192 | + cudaStream_t stream; |
| 193 | + ERROR_CHECK_CUDA(cudaStreamCreate(&stream)); |
| 194 | + std::cout << "--- Single event, synchronous wait for completion ---\n"; |
| 195 | + log() << "main Launching algorithm..." << std::endl; |
| 196 | + auto status = sync_wait(scheduler, reconstruct(stream, "main")); |
| 197 | + log() << "main Final status of algorithm " << status << "" << std::endl; |
| 198 | + ERROR_CHECK_CUDA(cudaStreamDestroy(stream)); |
| 199 | + } |
| 200 | + |
| 201 | + { |
| 202 | + std::cout << "--- Multiple events, wait for all to complete ---\n"; |
| 203 | + |
| 204 | + auto streams = std::vector<cudaStream_t>(2); |
| 205 | + auto status = std::vector<tool::StatusCode>(streams.size()); |
| 206 | + for (auto& stream : streams) { |
| 207 | + ERROR_CHECK_CUDA(cudaStreamCreate(&stream)); |
| 208 | + } |
| 209 | + |
| 210 | + auto scope = counting_scope{}; |
| 211 | + log() << "main Launching algorithms..." << std::endl; |
| 212 | + |
| 213 | + auto payload = [](std::vector<cudaStream_t> streams, |
| 214 | + std::vector<tool::StatusCode>& statuses, |
| 215 | + int i) -> tool::Task<void> { |
| 216 | + const auto name = std::format("event{}:main", i); |
| 217 | + auto& stream = streams.at(i); |
| 218 | + auto& status = statuses.at(i); |
| 219 | + status = co_await reconstruct(stream, name); |
| 220 | + }; |
| 221 | + for (std::size_t i = 0; i < streams.size(); ++i) { |
| 222 | + scope.spawn(scheduler, payload(streams, status, i)); |
| 223 | + } |
| 224 | + scope.join(); |
| 225 | + |
| 226 | + for (auto& stream : streams) { |
| 227 | + ERROR_CHECK_CUDA(cudaStreamDestroy(stream)); |
| 228 | + } |
| 229 | + |
| 230 | + log() << "main All algorithms completed. Final statuses: "; |
| 231 | + for (auto s : status) { |
| 232 | + std::cout << s << ' '; |
| 233 | + } |
| 234 | + std::cout << std::endl; |
| 235 | + } |
| 236 | + return 0; |
| 237 | +} |
0 commit comments