Skip to content

Commit b336809

Browse files
m-filam-fila
authored
add reco example with capy (#51)
1 parent 95e16de commit b336809

4 files changed

Lines changed: 286 additions & 2 deletions

File tree

‎examples/CMakeLists.txt‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,4 +73,8 @@ if(BUILD_CUDA AND BUILD_TBB)
7373
add_executable(exec_reco_stdexec exec_reco.cpp)
7474
target_link_libraries(exec_reco_stdexec PRIVATE CoroutineTests CUDA::cudart TBB::tbb STDEXEC::stdexec)
7575
endif()
76+
if(BUILD_CAPY)
77+
add_executable(capy_reco capy_reco.cpp)
78+
target_link_libraries(capy_reco PRIVATE CoroutineTests CUDA::cudart TBB::tbb Boost::capy)
79+
endif()
7680
endif()

‎examples/README.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,6 @@ This is a variant of ["Exec tbb" example](#exec-tbb) using Boost.Capy and IoAwai
158158

159159
## Reconstruction
160160

161-
Link: [alien_reco.cpp](alien_reco.cpp) and [exec_reco.cpp](exec_reco.cpp)
161+
Link: [alien_reco.cpp](alien_reco.cpp), [exec_reco.cpp](exec_reco.cpp), [capy_reco.cpp](capy_reco.cpp)
162162

163-
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.
163+
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)), in `exec_reco_stdexec` the application is implemented with C++26 execution, in `capy_reco` the implementation uses Boost.Capy and IoAwaitables protocol ([p4003](https://www.open-std.org/JTC1/SC22/WG21/docs/papers/2026/p4003r0.pdf)).

‎examples/capy_reco.cpp‎

Lines changed: 241 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,241 @@
1+
#include <cuda_runtime_api.h>
2+
#include <tbb/task_arena.h>
3+
4+
#include <boost/capy.hpp>
5+
#include <condition_variable>
6+
#include <cstddef>
7+
#include <iostream>
8+
#include <mutex>
9+
#include <string_view>
10+
11+
#include "capy_stream_await.hpp" // StreamIoAwaitable
12+
#include "capy_task_arena_executor.hpp" // TaskArenaExecutor
13+
#include "logging_utils.hpp" // log, format_name
14+
#include "statuscode.hpp" // StatusCodeImpl
15+
16+
namespace tools {
17+
struct Tag {
18+
static constexpr const char* name = "tools";
19+
};
20+
using StatusCode = StatusCodeImpl<Tag>;
21+
} // namespace tools
22+
23+
#define ERROR_CHECK_CUDA(EXP) \
24+
do { \
25+
cudaError_t errorCode = EXP; \
26+
if (errorCode != cudaSuccess) { \
27+
throw std::runtime_error( \
28+
std::format("CUDA error at {}:{}: {}", __FILE__, __LINE__, \
29+
cudaGetErrorString(errorCode))); \
30+
} \
31+
} while (false)
32+
33+
template <typename T>
34+
struct DeviceBuffer {
35+
T* ptr = nullptr;
36+
std::size_t size = 0;
37+
};
38+
39+
boost::capy::task<DeviceBuffer<int>> clusterization(DeviceBuffer<int> cells,
40+
cudaStream_t stream,
41+
std::string_view parent) {
42+
43+
const auto self = format_name(parent, "clusterization");
44+
log(self) << "Starting clusterization" << std::endl;
45+
46+
const auto nCells = static_cast<int>(cells.size);
47+
48+
// Copy cells back to host to count non-zero entries
49+
auto h_cells = std::vector<int>(nCells);
50+
51+
ERROR_CHECK_CUDA(cudaMemcpyAsync(h_cells.data(), cells.ptr,
52+
nCells * sizeof(int),
53+
cudaMemcpyDeviceToHost, stream));
54+
55+
ERROR_CHECK_CUDA(co_await StreamIoAwaitable{stream});
56+
57+
auto nClusters = 0;
58+
for (auto v : h_cells)
59+
if (v != 0)
60+
++nClusters;
61+
62+
log(self) << "Found " << nClusters << " clusters" << std::endl;
63+
64+
// Allocate clusters of appropiate size on device
65+
int* d_clusters = nullptr;
66+
ERROR_CHECK_CUDA(cudaMallocAsync(reinterpret_cast<void**>(&d_clusters),
67+
nClusters * sizeof(int), stream));
68+
69+
// Write some dummy data to the clusters buffer to simulate work
70+
ERROR_CHECK_CUDA(
71+
cudaMemsetAsync(d_clusters, 0, nClusters * sizeof(int), stream));
72+
ERROR_CHECK_CUDA(
73+
cudaMemsetAsync(d_clusters, 1, nClusters / 2 * sizeof(int), stream));
74+
75+
co_return DeviceBuffer<int>{d_clusters,
76+
static_cast<std::size_t>(nClusters)};
77+
}
78+
79+
boost::capy::task<DeviceBuffer<int>> seeding(DeviceBuffer<int> clusters,
80+
cudaStream_t stream,
81+
std::string_view parent) {
82+
83+
const auto self = format_name(parent, "seeding");
84+
log(self) << "Starting seeding" << std::endl;
85+
86+
const auto nClusters = static_cast<int>(clusters.size);
87+
88+
// Copy clusters to host to count non-zero entries
89+
auto h_clusters = std::vector<int>(nClusters);
90+
91+
ERROR_CHECK_CUDA(cudaMemcpyAsync(h_clusters.data(), clusters.ptr,
92+
nClusters * sizeof(int),
93+
cudaMemcpyDeviceToHost, stream));
94+
95+
ERROR_CHECK_CUDA(co_await StreamIoAwaitable{stream});
96+
97+
int nSeeds = 0;
98+
for (auto v : h_clusters)
99+
if (v != 0)
100+
++nSeeds;
101+
102+
log(self) << "Found " << nSeeds << " seeds" << std::endl;
103+
104+
// Allocate clusters of appropiate size on device
105+
int* d_seeds = nullptr;
106+
ERROR_CHECK_CUDA(cudaMallocAsync(reinterpret_cast<void**>(&d_seeds),
107+
nSeeds * sizeof(int), stream));
108+
109+
// Write some dummy data to the clusters buffer to simulate work
110+
ERROR_CHECK_CUDA(cudaMemsetAsync(d_seeds, 0, nSeeds * sizeof(int), stream));
111+
ERROR_CHECK_CUDA(
112+
cudaMemsetAsync(d_seeds, 1, nSeeds / 2 * sizeof(int), stream));
113+
114+
co_return DeviceBuffer<int>{d_seeds, static_cast<std::size_t>(nSeeds)};
115+
}
116+
117+
boost::capy::task<tools::StatusCode> reconstruct(cudaStream_t stream,
118+
std::string parent) {
119+
const auto self = format_name(parent, "reconstruction");
120+
log(self) << "Starting reconstruction" << std::endl;
121+
122+
// Allocate some dummy input data on the device
123+
auto cells = DeviceBuffer<int>{nullptr, 1000};
124+
ERROR_CHECK_CUDA(cudaMallocAsync(reinterpret_cast<void**>(&cells.ptr),
125+
cells.size * sizeof(int), stream));
126+
ERROR_CHECK_CUDA(
127+
cudaMemsetAsync(cells.ptr, 1, cells.size * sizeof(int), stream));
128+
129+
// Run the clusterization and seeding steps
130+
auto clusters = co_await clusterization(cells, stream, self);
131+
auto seeds = co_await seeding(clusters, stream, self);
132+
133+
// Cleanup
134+
ERROR_CHECK_CUDA(cudaFreeAsync(cells.ptr, stream));
135+
ERROR_CHECK_CUDA(cudaFreeAsync(clusters.ptr, stream));
136+
ERROR_CHECK_CUDA(cudaFreeAsync(seeds.ptr, stream));
137+
138+
ERROR_CHECK_CUDA(co_await StreamIoAwaitable{stream});
139+
140+
log(self) << "Finishing reconstruction" << std::endl;
141+
co_return tools::StatusCode::SUCCESS;
142+
}
143+
144+
int main() {
145+
int deviceCount = 0;
146+
auto error_id = cudaGetDeviceCount(&deviceCount);
147+
148+
if (error_id != cudaSuccess) {
149+
std::cout << "cudaGetDeviceCount returned "
150+
<< static_cast<int>(error_id) << "\n"
151+
<< cudaGetErrorString(error_id) << "\n";
152+
return EXIT_FAILURE;
153+
}
154+
155+
if (deviceCount == 0) {
156+
std::cout << "No CUDA devices found.\n";
157+
return EXIT_FAILURE;
158+
}
159+
160+
log() << "main Starting" << std::endl;
161+
162+
auto task_arena = tbb::task_arena{2};
163+
auto context = TaskArenaContext(task_arena);
164+
auto executor = TaskArenaExecutor(context);
165+
166+
{
167+
std::cout << "--- Single event, synchronous wait for completion ---\n";
168+
cudaStream_t stream;
169+
ERROR_CHECK_CUDA(cudaStreamCreate(&stream));
170+
auto condition = std::condition_variable();
171+
auto mutex = std::mutex();
172+
auto final_result = tools::StatusCode{};
173+
auto result_handler = [&condition, &mutex,
174+
&final_result](tools::StatusCode code) {
175+
{
176+
std::lock_guard lock(mutex);
177+
final_result = code;
178+
}
179+
condition.notify_one();
180+
};
181+
182+
log() << "main Launching algorithm..." << std::endl;
183+
boost::capy::run_async(executor,
184+
result_handler)(reconstruct(stream, "main"));
185+
{
186+
auto lock = std::unique_lock(mutex);
187+
condition.wait(lock, [&final_result]() {
188+
return final_result != tools::StatusCode::UNDEFINED;
189+
});
190+
}
191+
log() << "main Final status of algorithm " << final_result << ""
192+
<< std::endl;
193+
ERROR_CHECK_CUDA(cudaStreamDestroy(stream));
194+
}
195+
196+
{
197+
std::cout << "--- Multiple events, wait for all to complete ---\n ";
198+
199+
auto streams = std::vector<cudaStream_t>(2);
200+
auto status = std::vector<tools::StatusCode>(streams.size());
201+
for (auto& stream : streams) {
202+
ERROR_CHECK_CUDA(cudaStreamCreate(&stream));
203+
}
204+
205+
auto condition = std::condition_variable();
206+
auto mutex = std::mutex();
207+
auto counter = streams.size();
208+
log() << "main Launching algorithms..." << std::endl;
209+
210+
for (std::size_t i = 0; i < streams.size(); ++i) {
211+
auto result_handler = [&condition, &mutex, &status, &counter,
212+
i](tools::StatusCode code) {
213+
auto done = false;
214+
{
215+
std::lock_guard lock(mutex);
216+
status.at(i) = code;
217+
done = (--counter == 0);
218+
}
219+
if (done) {
220+
condition.notify_one();
221+
}
222+
};
223+
boost::capy::run_async(executor, result_handler)(
224+
reconstruct(streams.at(i), "event" + std::to_string(i)));
225+
}
226+
{
227+
auto lock = std::unique_lock(mutex);
228+
condition.wait(lock, [&counter]() { return counter == 0; });
229+
}
230+
for (auto& stream : streams) {
231+
ERROR_CHECK_CUDA(cudaStreamDestroy(stream));
232+
}
233+
234+
log() << "main All algorithms completed. Final statuses: ";
235+
for (auto s : status) {
236+
std::cout << s << ' ';
237+
}
238+
std::cout << std::endl;
239+
}
240+
return 0;
241+
}

‎examples/capy_stream_await.hpp‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
#pragma once
2+
#include <cuda_runtime_api.h>
3+
4+
#include <boost/capy.hpp>
5+
#include <coroutine>
6+
7+
class StreamIoAwaitable {
8+
public:
9+
StreamIoAwaitable(cudaStream_t stream) : m_stream(stream) {}
10+
11+
bool await_ready() const noexcept { return false; }
12+
13+
void await_suspend(std::coroutine_handle<> handle,
14+
boost::capy::io_env const* env) noexcept {
15+
m_context.handle = handle;
16+
m_context.env = env;
17+
m_error = cudaLaunchHostFunc(m_stream, resumption_callback, &m_context);
18+
// If the callback couldn't be registered, we need to reschedule the
19+
// coroutine immediately to avoid deadlock.
20+
if (m_error != cudaSuccess) {
21+
resumption_callback(&m_context);
22+
}
23+
}
24+
cudaError_t await_resume() const noexcept { return m_error; }
25+
26+
private:
27+
struct context {
28+
std::coroutine_handle<> handle;
29+
boost::capy::io_env const* env;
30+
};
31+
cudaStream_t m_stream;
32+
cudaError_t m_error = cudaSuccess;
33+
context m_context;
34+
35+
static void resumption_callback(void* userData) {
36+
auto* ctx = static_cast<context*>(userData);
37+
ctx->env->executor.post(ctx->handle);
38+
}
39+
};

0 commit comments

Comments
 (0)