Skip to content

Commit e2ac8af

Browse files
hkaiserhkaiser
authored
Merge pull request #7330 from guptapratykshh/feat/cuda-env
async_cuda: use member dispatch for executor entry points
2 parents 04431bf + 2088956 commit e2ac8af

4 files changed

Lines changed: 64 additions & 29 deletions

File tree

‎libs/core/async_cuda/include/hpx/async_cuda/cublas_executor.hpp‎

Lines changed: 23 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -134,11 +134,27 @@ namespace hpx::cuda::experimental {
134134
// -------------------------------------------------------------------------
135135
// OneWay Execution
136136
// -------------------------------------------------------------------------
137+
template <typename F, typename... Ts>
138+
void post(F&& f, Ts&&... ts) const
139+
{
140+
post_impl(HPX_FORWARD(F, f), HPX_FORWARD(Ts, ts)...);
141+
}
142+
143+
// -------------------------------------------------------------------------
144+
// TwoWay Execution
145+
// -------------------------------------------------------------------------
146+
template <typename F, typename... Ts>
147+
decltype(auto) async_execute(F&& f, Ts&&... ts) const
148+
{
149+
return async_impl(HPX_FORWARD(F, f), HPX_FORWARD(Ts, ts)...);
150+
}
151+
152+
protected:
137153
// This is a simple wrapper for any cublas call, pass in the same arguments
138154
// that you would use for a cublas call except the cublas handle which is omitted
139155
// as the wrapper will supply that for you
140156
template <typename R, typename... Params, typename... Args>
141-
std::enable_if_t<std::is_same_v<cublasStatus_t, R>, R> post(
157+
std::enable_if_t<std::is_same_v<cublasStatus_t, R>, R> post_impl(
142158
R (*cublas_function)(Params...), Args&&... args) const
143159
{
144160
// make sure we run on the correct device
@@ -159,30 +175,20 @@ namespace hpx::cuda::experimental {
159175
// forward a cuda function through to the cuda executor base class
160176
// (we permit the use of a cublas executor for cuda calls)
161177
template <typename R, typename... Params, typename... Args>
162-
inline std::enable_if_t<std::is_same_v<cudaError_t, R>> post(
178+
inline std::enable_if_t<std::is_same_v<cudaError_t, R>> post_impl(
163179
R (*cuda_function)(Params...), Args&&... args) const
164180
{
165-
return cuda_executor::post(
181+
return cuda_executor::post_impl(
166182
cuda_function, HPX_FORWARD(Args, args)...);
167183
}
168184

169-
// -------------------------------------------------------------------------
170-
// TwoWay Execution
171-
// -------------------------------------------------------------------------
172-
template <typename F, typename... Ts>
173-
decltype(auto) async_execute(F&& f, Ts&&... ts) const
174-
{
175-
return async(HPX_FORWARD(F, f), HPX_FORWARD(Ts, ts)...);
176-
}
177-
178-
protected:
179185
// -------------------------------------------------------------------------
180186
// launch a cuBlas function and return a future that will become ready
181187
// when the task completes, this allows integration of GPU kernels with
182188
// hpx::futures and the tasking DAG.
183189
template <typename R, typename... Params, typename... Args>
184-
hpx::future<std::enable_if_t<std::is_same_v<cublasStatus_t, R>>> async(
185-
R (*cublas_function)(Params...), Args&&... args) const
190+
hpx::future<std::enable_if_t<std::is_same_v<cublasStatus_t, R>>>
191+
async_impl(R (*cublas_function)(Params...), Args&&... args) const
186192
{
187193
return hpx::detail::try_catch_exception_ptr(
188194
[&]() {
@@ -208,9 +214,9 @@ namespace hpx::cuda::experimental {
208214
// forward a cuda function through to the cuda executor base class
209215
template <typename R, typename... Params, typename... Args>
210216
inline hpx::future<std::enable_if_t<std::is_same_v<cudaError_t, R>>>
211-
async(R (*cuda_function)(Params...), Args&&... args) const
217+
async_impl(R (*cuda_function)(Params...), Args&&... args) const
212218
{
213-
return cuda_executor::async(
219+
return cuda_executor::async_impl(
214220
cuda_function, HPX_FORWARD(Args, args)...);
215221
}
216222

‎libs/core/async_cuda/include/hpx/async_cuda/cuda_executor.hpp‎

Lines changed: 23 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -125,32 +125,44 @@ namespace hpx::cuda::experimental {
125125

126126
// -------------------------------------------------------------------------
127127
// OneWay Execution
128-
template <typename R, typename... Params, typename... Args>
129-
void post(R (*cuda_function)(Params...), Args&&... args) const
128+
template <typename F, typename... Ts>
129+
void post(F&& f, Ts&&... ts) const
130130
{
131-
// make sure we run on the correct device
132-
check_cuda_error(cudaSetDevice(device_));
133-
134-
// insert the stream handle in the arg list and call the cuda function
135-
detail::dispatch_helper<R, Params...> helper{};
136-
helper(cuda_function, HPX_FORWARD(Args, args)..., stream_);
131+
post_impl(HPX_FORWARD(F, f), HPX_FORWARD(Ts, ts)...);
137132
}
138133

139134
// -------------------------------------------------------------------------
140135
// TwoWay Execution
141136
template <typename F, typename... Ts>
142137
decltype(auto) async_execute(F&& f, Ts&&... ts) const
143138
{
144-
return async(HPX_FORWARD(F, f), HPX_FORWARD(Ts, ts)...);
139+
return async_impl(HPX_FORWARD(F, f), HPX_FORWARD(Ts, ts)...);
140+
}
141+
142+
protected:
143+
// -------------------------------------------------------------------------
144+
// launch a kernel on our stream without returning a future
145+
// Errors are reported through check_cuda_error, typically via the
146+
// CUDA call return value (for example cudaError_t).
147+
// Throws cuda_exception if the async launch fails.
148+
template <typename R, typename... Params, typename... Args>
149+
void post_impl(R (*cuda_function)(Params...), Args&&... args) const
150+
{
151+
// make sure we run on the correct device
152+
check_cuda_error(cudaSetDevice(device_));
153+
154+
// insert the stream handle in the arg list and call the cuda function
155+
detail::dispatch_helper<R, Params...> helper{};
156+
helper(cuda_function, HPX_FORWARD(Args, args)..., stream_);
145157
}
146158

147159
// -------------------------------------------------------------------------
148160
// launch a kernel on our stream and return a future that will become ready
149-
// when the task completes, this allows integregration of GPU kernels with
161+
// when the task completes, this allows integration of GPU kernels with
150162
// hpx::futures and the tasking DAG.
151163
// Puts a cuda_exception in the future if the async launch fails.
152164
template <typename R, typename... Params, typename... Args>
153-
hpx::future<void> async(
165+
hpx::future<void> async_impl(
154166
R (*cuda_kernel)(Params...), Args&&... args) const
155167
{
156168
return hpx::detail::try_catch_exception_ptr(

‎libs/core/async_cuda/include/hpx/async_cuda/transform_stream.hpp‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -327,7 +327,8 @@ namespace hpx::cuda::experimental {
327327
}
328328
// clang-format on
329329

330-
constexpr auto get_env() const noexcept
330+
constexpr decltype(auto) get_env() const
331+
noexcept(noexcept(hpx::execution::experimental::get_env(s)))
331332
{
332333
return hpx::execution::experimental::get_env(s);
333334
}

‎libs/core/async_cuda/tests/unit/transform_stream.cu‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
#include <atomic>
1919
#include <cstddef>
2020
#include <iostream>
21+
#include <type_traits>
2122
#include <utility>
2223

2324
__global__ void dummy_kernel() {}
@@ -115,6 +116,21 @@ int hpx_main()
115116
namespace ex = ::hpx::execution::experimental;
116117
namespace tt = hpx::this_thread::experimental;
117118

119+
{
120+
auto s = ex::schedule(ex::thread_pool_scheduler{});
121+
using sender_type = decltype(s);
122+
using original_scheduler_type =
123+
std::decay_t<decltype(ex::get_completion_scheduler<ex::set_value_t>(
124+
ex::get_env(std::declval<sender_type&>())))>;
125+
126+
auto transformed = cu::transform_stream(std::move(s), dummy{});
127+
using scheduler_type =
128+
std::decay_t<decltype(ex::get_completion_scheduler<ex::set_value_t>(
129+
ex::get_env(transformed)))>;
130+
131+
static_assert(std::is_same_v<scheduler_type, original_scheduler_type>);
132+
}
133+
118134
cu::enable_user_polling p;
119135

120136
// Only stream transform

0 commit comments

Comments
 (0)