Skip to content

Commit aa2ef3a

Browse files
ericnieblerericniebler
committed
help algorithm dispatch in nvexec tests
1 parent a2e6e7d commit aa2ef3a

3 files changed

Lines changed: 45 additions & 34 deletions

File tree

‎include/nvexec/stream/common.cuh‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -439,7 +439,8 @@ namespace nv::execution
439439
return get_stream_provider(env)->own_stream_.value();
440440
}
441441

442-
STDEXEC_ATTRIBUTE(host, device) auto operator()() const noexcept
442+
STDEXEC_ATTRIBUTE(host, device)
443+
auto operator()() const noexcept
443444
{
444445
return STDEXEC::read_env(*this);
445446
}

‎test/nvexec/bulk.cpp‎

Lines changed: 27 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,12 @@
1-
#include <stdexec/execution.hpp>
21
#include <test_common/catch2.hpp>
32

3+
#include <stdexec/execution.hpp>
4+
5+
#include <exec/env.hpp>
6+
7+
#include <nvexec/stream_context.cuh>
8+
49
#include "common.cuh"
5-
#include "nvexec/stream_context.cuh"
610

711
#include <cuda/std/span>
812

@@ -26,7 +30,7 @@ namespace
2630
nvexec::stream_context stream_ctx{};
2731
auto const snd = ex::schedule(stream_ctx.get_scheduler()) | ex::bulk(ex::par, 1, [](int) {});
2832

29-
REQUIRE(STDEXEC::sync_wait(snd).has_value());
33+
REQUIRE(ex::sync_wait(snd).has_value());
3034
}
3135

3236
TEST_CASE("nvexec bulk executes on GPU", "[cuda][stream][adaptors][bulk]")
@@ -46,7 +50,7 @@ namespace
4650
flags.set(idx);
4751
}
4852
});
49-
STDEXEC::sync_wait(std::move(snd));
53+
ex::sync_wait(std::move(snd));
5054

5155
REQUIRE(flags_storage.all_set_once());
5256
}
@@ -71,7 +75,7 @@ namespace
7175
}
7276
}
7377
});
74-
STDEXEC::sync_wait(std::move(snd));
78+
ex::sync_wait(std::move(snd));
7579

7680
REQUIRE(flags_storage.all_set_once());
7781
}
@@ -96,7 +100,7 @@ namespace
96100
}
97101
}
98102
});
99-
auto const [i, d] = STDEXEC::sync_wait(std::move(snd)).value();
103+
auto const [i, d] = ex::sync_wait(std::move(snd)).value();
100104

101105
REQUIRE(flags_storage.all_set_once());
102106
REQUIRE(i == 42);
@@ -122,7 +126,7 @@ namespace
122126
flags.set(idx);
123127
}
124128
});
125-
[[maybe_unused]] auto [flags_actual] = STDEXEC::sync_wait(std::move(snd)).value();
129+
[[maybe_unused]] auto [flags_actual] = ex::sync_wait(std::move(snd)).value();
126130

127131
REQUIRE(flags_storage.all_set_once());
128132
}
@@ -152,7 +156,7 @@ namespace
152156
flags.set(2);
153157
}
154158
});
155-
STDEXEC::sync_wait(std::move(snd));
159+
ex::sync_wait(std::move(snd));
156160

157161
REQUIRE(flags_storage.all_set_once());
158162
}
@@ -183,7 +187,7 @@ namespace
183187
flags.set(idx);
184188
}
185189
});
186-
STDEXEC::sync_wait(std::move(snd));
190+
ex::sync_wait(std::move(snd));
187191

188192
REQUIRE(flags_storage.all_set_once());
189193
}
@@ -205,7 +209,7 @@ namespace
205209
flags.set(idx);
206210
}
207211
});
208-
STDEXEC::sync_wait(std::move(snd)).value();
212+
ex::sync_wait(std::move(snd)).value();
209213

210214
REQUIRE(flags_storage.all_set_once());
211215
}
@@ -220,19 +224,19 @@ namespace
220224
int const nelems = 10;
221225
cudaMallocManaged(&inout, nelems * sizeof(double));
222226

223-
auto task = STDEXEC::just(cuda::std::span<double>{inout, nelems})
224-
| STDEXEC::continues_on(ctx.get_scheduler())
225-
| STDEXEC::bulk(ex::par,
226-
nelems,
227-
[](std::size_t i, cuda::std::span<double> out)
228-
{ out[i] = (double) i; })
229-
| STDEXEC::let_value([](cuda::std::span<double> out) { return STDEXEC::just(out); })
230-
| STDEXEC::bulk(ex::par,
231-
nelems,
232-
[](std::size_t i, cuda::std::span<double> out)
233-
{ out[i] = 2.0 * out[i]; });
234-
235-
STDEXEC::sync_wait(std::move(task)).value();
227+
auto task = ex::just(cuda::std::span<double>{inout, nelems})
228+
| ex::continues_on(ctx.get_scheduler())
229+
| ex::bulk(ex::par,
230+
nelems,
231+
[](std::size_t i, cuda::std::span<double> out) { out[i] = (double) i; })
232+
| ex::let_value([](cuda::std::span<double> out) { return ex::just(out); })
233+
| exec::write_attrs(
234+
ex::prop{ex::get_completion_scheduler<ex::set_value_t>, ctx.get_scheduler()})
235+
| ex::bulk(ex::par,
236+
nelems,
237+
[](std::size_t i, cuda::std::span<double> out) { out[i] = 2.0 * out[i]; });
238+
239+
ex::sync_wait(std::move(task)).value();
236240

237241
for (int i = 0; i < nelems; ++i)
238242
{

‎test/nvexec/let_value.cpp‎

Lines changed: 16 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,12 @@
1-
#include <stdexec/execution.hpp>
21
#include <test_common/catch2.hpp>
32

3+
#include <stdexec/execution.hpp>
4+
5+
#include <exec/env.hpp>
6+
7+
#include <nvexec/stream_context.cuh>
8+
49
#include "common.cuh"
5-
#include "nvexec/stream_context.cuh"
610

711
namespace ex = STDEXEC;
812

@@ -36,7 +40,7 @@ namespace
3640
}
3741
return ex::just();
3842
});
39-
STDEXEC::sync_wait(std::move(snd));
43+
ex::sync_wait(std::move(snd));
4044

4145
REQUIRE(flags_storage.all_set_once());
4246
}
@@ -61,7 +65,7 @@ namespace
6165
}
6266
return ex::just();
6367
});
64-
STDEXEC::sync_wait(std::move(snd));
68+
ex::sync_wait(std::move(snd));
6569

6670
REQUIRE(flags_storage.all_set_once());
6771
}
@@ -87,7 +91,7 @@ namespace
8791
}
8892
return ex::just();
8993
});
90-
STDEXEC::sync_wait(std::move(snd));
94+
ex::sync_wait(std::move(snd));
9195

9296
REQUIRE(flags_storage.all_set_once());
9397
}
@@ -98,7 +102,7 @@ namespace
98102

99103
auto snd = ex::schedule(stream_ctx.get_scheduler())
100104
| ex::let_value([=]() { return ex::just(is_on_gpu()); });
101-
auto const [result] = STDEXEC::sync_wait(std::move(snd)).value();
105+
auto const [result] = ex::sync_wait(std::move(snd)).value();
102106

103107
REQUIRE(result == 1);
104108
}
@@ -130,7 +134,7 @@ namespace
130134
flags.set(1);
131135
}
132136
});
133-
STDEXEC::sync_wait(std::move(snd));
137+
ex::sync_wait(std::move(snd));
134138

135139
REQUIRE(flags_storage.all_set_once());
136140
}
@@ -153,7 +157,7 @@ namespace
153157

154158
return ex::schedule(sch);
155159
});
156-
STDEXEC::sync_wait(std::move(snd));
160+
ex::sync_wait(std::move(snd));
157161

158162
REQUIRE(flags_storage.all_set_once());
159163
}
@@ -165,7 +169,9 @@ namespace
165169
flags_storage_t flags_storage{};
166170
auto flags = flags_storage.get();
167171

168-
auto snd = ex::schedule(sch) | ex::let_value([] { return nvexec::get_stream(); })
172+
auto snd = ex::schedule(sch) //
173+
| ex::let_value([] { return nvexec::get_stream(); }) //
174+
| exec::write_attrs(ex::prop{ex::get_completion_scheduler<ex::set_value_t>, sch}) //
169175
| ex::then(
170176
[flags](cudaStream_t stream)
171177
{
@@ -175,7 +181,7 @@ namespace
175181
}
176182
return stream;
177183
});
178-
auto [stream] = STDEXEC::sync_wait(std::move(snd)).value();
184+
auto [stream] = ex::sync_wait(std::move(snd)).value();
179185
static_assert(std::same_as<decltype(+stream), cudaStream_t>);
180186

181187
REQUIRE(flags_storage.all_set_once());

0 commit comments

Comments
 (0)