exec_policy.hpp
Go to the documentation of this file.
1 /*
2  * Copyright (c) 2020-2021, NVIDIA CORPORATION.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
22 #pragma once
23 
24 #include <rmm/cuda_stream_view.hpp>
25 #include <rmm/mr/device/thrust_allocator_adaptor.hpp>
26 
27 #include <thrust/system/cuda/execution_policy.h>
28 #include <thrust/version.h>
29 
30 namespace rmm {
31 
32 using thrust_exec_policy_t =
33  thrust::detail::execute_with_allocator<rmm::mr::thrust_allocator<char>,
34  thrust::cuda_cub::execute_on_stream_base>;
35 
40 class exec_policy : public thrust_exec_policy_t {
41  public:
42  explicit exec_policy(cuda_stream_view stream = cuda_stream_default,
43  rmm::mr::device_memory_resource* mr = mr::get_current_device_resource())
44  : thrust_exec_policy_t(
45  thrust::cuda::par(rmm::mr::thrust_allocator<char>(stream, mr)).on(stream.value()))
46  {
47  }
48 };
49 
50 #if THRUST_VERSION >= 101600
51 
52 using thrust_exec_policy_nosync_t =
53  thrust::detail::execute_with_allocator<rmm::mr::thrust_allocator<char>,
54  thrust::cuda_cub::execute_on_stream_nosync_base>;
61 class exec_policy_nosync : public thrust_exec_policy_nosync_t {
62  public:
63  explicit exec_policy_nosync(
64  cuda_stream_view stream = cuda_stream_default,
65  rmm::mr::device_memory_resource* mr = mr::get_current_device_resource())
66  : thrust_exec_policy_nosync_t(
67  thrust::cuda::par_nosync(rmm::mr::thrust_allocator<char>(stream, mr)).on(stream.value()))
68  {
69  }
70 };
71 
72 #else
73 
74 using thrust_exec_policy_nosync_t = thrust_exec_policy_t;
75 using exec_policy_nosync = exec_policy;
76 
77 #endif
78 
79 } // namespace rmm
rmm::mr::thrust_allocator
An allocator compatible with Thrust containers and algorithms using a device_memory_resource for memo...
Definition: thrust_allocator_adaptor.hpp:38
rmm::cuda_stream_view
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:34
rmm::mr::device_memory_resource
Base class for all libcudf device memory allocation.
Definition: device_memory_resource.hpp:82
rmm::exec_policy
Helper class usable as a Thrust CUDA execution policy that uses RMM for temporary memory allocation o...
Definition: exec_policy.hpp:40