Skip to content

Commit bd3ac04

Browse files
nicolasnoblenicolasnoble
committed
Starting to work on coroutines.
1 parent 6cc1c11 commit bd3ac04

3 files changed

Lines changed: 232 additions & 0 deletions

File tree

‎src/CMakeLists.txt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -194,6 +194,7 @@ set(
194194
periodic_concurrency_worker.h
195195
thread_config.h
196196
cuda_runtime_library_manager.h
197+
coroutines.h
197198
)
198199

199200
add_executable(
@@ -338,6 +339,7 @@ add_executable(
338339
test_ctx_id_tracker.cc
339340
test_profile_data_collector.cc
340341
test_profile_data_exporter.cc
342+
test_coroutines.cc
341343
${TEST_HTTP_CLIENT}
342344
)
343345

@@ -361,6 +363,7 @@ target_compile_options(
361363
perf_analyzer_unit_tests
362364
PRIVATE
363365
-Wno-write-strings
366+
-fcoroutines
364367
)
365368

366369
target_link_libraries(

‎src/coroutines.h‎

Lines changed: 173 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,173 @@
1+
// Copyright 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
//
3+
// Redistribution and use in source and binary forms, with or without
4+
// modification, are permitted provided that the following conditions
5+
// are met:
6+
// * Redistributions of source code must retain the above copyright
7+
// notice, this list of conditions and the following disclaimer.
8+
// * Redistributions in binary form must reproduce the above copyright
9+
// notice, this list of conditions and the following disclaimer in the
10+
// documentation and/or other materials provided with the distribution.
11+
// * Neither the name of NVIDIA CORPORATION nor the names of its
12+
// contributors may be used to endorse or promote products derived
13+
// from this software without specific prior written permission.
14+
//
15+
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
16+
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
17+
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
18+
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
19+
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
20+
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
21+
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
22+
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
23+
// OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
24+
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
25+
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
26+
#pragma once
27+
28+
#include <coroutine>
29+
#include <type_traits>
30+
#include <utility>
31+
32+
namespace triton::perfanalyzer {
33+
34+
template <typename T = void>
35+
class Coroutine {
36+
public:
37+
struct Empty {};
38+
typedef
39+
typename std::conditional<std::is_void<T>::value, Empty, T>::type SafeT;
40+
41+
Coroutine() = default;
42+
Coroutine(Coroutine&& other) = default;
43+
Coroutine& operator=(Coroutine&& other) = default;
44+
Coroutine(Coroutine const&) = delete;
45+
Coroutine& operator=(Coroutine const&) = delete;
46+
47+
class Awaiter {
48+
public:
49+
Awaiter(Awaiter&& other) = default;
50+
Awaiter& operator=(Awaiter&& other) = default;
51+
Awaiter(Awaiter const&) = default;
52+
Awaiter& operator=(Awaiter const&) = default;
53+
constexpr bool await_ready() const noexcept
54+
{
55+
bool ret = coroutine_->earlyResume_;
56+
coroutine_->earlyResume_ = false;
57+
return ret;
58+
}
59+
constexpr void await_suspend(std::coroutine_handle<> h)
60+
{
61+
coroutine_->suspended_ = true;
62+
}
63+
constexpr void await_resume() const noexcept {}
64+
65+
private:
66+
Awaiter(Coroutine* coroutine) : coroutine_(coroutine) {}
67+
Coroutine* coroutine_;
68+
friend struct Coroutine;
69+
};
70+
71+
Awaiter awaiter() { return Awaiter(this); }
72+
73+
void resume()
74+
{
75+
if (!handle_)
76+
return;
77+
if (!suspended_) {
78+
earlyResume_ = true;
79+
return;
80+
}
81+
suspended_ = false;
82+
handle_.resume();
83+
}
84+
85+
bool done()
86+
{
87+
if (!handle_)
88+
return true;
89+
bool isDone = handle_.done();
90+
if (isDone) {
91+
if constexpr (!std::is_void<T>::value) {
92+
value_ = std::move(handle_.promise().value_);
93+
}
94+
handle_.destroy();
95+
handle_ = nullptr;
96+
}
97+
return isDone;
98+
}
99+
100+
const SafeT& value() const { return value_; }
101+
102+
private:
103+
struct PromiseVoid {
104+
Coroutine<> get_return_object()
105+
{
106+
return Coroutine<>{
107+
std::move(std::coroutine_handle<Promise>::from_promise(*this))};
108+
}
109+
std::suspend_always initial_suspend() { return {}; }
110+
std::suspend_always final_suspend() noexcept { return {}; }
111+
void unhandled_exception() {}
112+
void return_void()
113+
{
114+
auto awaitingCoroutine = awaitingCoroutine_;
115+
if (awaitingCoroutine) {
116+
__builtin_coro_resume(awaitingCoroutine);
117+
}
118+
}
119+
[[no_unique_address]] Empty value_;
120+
void* awaitingCoroutine_ = nullptr;
121+
};
122+
struct PromiseValue {
123+
Coroutine<T> get_return_object()
124+
{
125+
return Coroutine{
126+
std::move(std::coroutine_handle<Promise>::from_promise(*this))};
127+
}
128+
std::suspend_always initial_suspend() { return {}; }
129+
std::suspend_always final_suspend() noexcept { return {}; }
130+
void unhandled_exception() {}
131+
void return_value(T&& value)
132+
{
133+
value_ = std::move(value);
134+
auto awaitingCoroutine = awaitingCoroutine_;
135+
if (awaitingCoroutine) {
136+
__builtin_coro_resume(awaitingCoroutine);
137+
}
138+
}
139+
T value_;
140+
void* awaitingCoroutine_ = nullptr;
141+
};
142+
typedef typename std::conditional<
143+
std::is_void<T>::value, PromiseVoid, PromiseValue>::type Promise;
144+
Coroutine(std::coroutine_handle<Promise>&& handle)
145+
: handle_(std::move(handle))
146+
{
147+
}
148+
std::coroutine_handle<Promise> handle_;
149+
[[no_unique_address]] SafeT value_;
150+
void* awaitingCoroutine_ = nullptr;
151+
bool suspended_ = true;
152+
bool earlyResume_ = false;
153+
154+
public:
155+
using promise_type = Promise;
156+
157+
constexpr bool await_ready() { return handle_.done(); }
158+
template <typename U>
159+
constexpr void await_suspend(std::coroutine_handle<U> h)
160+
{
161+
auto& promise = handle_.promise();
162+
promise.awaitingCoroutine_ = h.address();
163+
resume();
164+
}
165+
constexpr SafeT await_resume()
166+
{
167+
SafeT value = std::move(handle_.promise().value_);
168+
handle_.destroy();
169+
return value;
170+
}
171+
};
172+
173+
} // namespace triton::perfanalyzer

‎src/test_coroutines.cc‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
// Copyright 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
//
3+
// Redistribution and use in source and binary forms, with or without
4+
// modification, are permitted provided that the following conditions
5+
// are met:
6+
// * Redistributions of source code must retain the above copyright
7+
// notice, this list of conditions and the following disclaimer.
8+
// * Redistributions in binary form must reproduce the above copyright
9+
// notice, this list of conditions and the following disclaimer in the
10+
// documentation and/or other materials provided with the distribution.
11+
// * Neither the name of NVIDIA CORPORATION nor the names of its
12+
// contributors may be used to endorse or promote products derived
13+
// from this software without specific prior written permission.
14+
//
15+
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
16+
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
17+
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
18+
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
19+
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
20+
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
21+
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
22+
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
23+
// OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
24+
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
25+
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
26+
27+
#include "coroutines.h"
28+
#include "doctest.h"
29+
30+
namespace triton::perfanalyzer {
31+
32+
Coroutine<int>
33+
CoroutineTest()
34+
{
35+
co_await std::suspend_always{};
36+
co_return 42;
37+
}
38+
39+
TEST_CASE("testing the Coroutine class")
40+
{
41+
auto coroutine = CoroutineTest();
42+
43+
unsigned rounds = 0;
44+
while (!coroutine.done()) {
45+
coroutine.resume();
46+
rounds++;
47+
}
48+
49+
auto result = coroutine.value();
50+
51+
CHECK(rounds == 2);
52+
CHECK(result == 42);
53+
CHECK(coroutine.done());
54+
}
55+
56+
} // namespace triton::perfanalyzer

0 commit comments

Comments
 (0)