24 #ifndef TILEDARRAY_CUDA_UM_VECTOR_H__INCLUDED
25 #define TILEDARRAY_CUDA_UM_VECTOR_H__INCLUDED
30 #ifdef TILEDARRAY_HAS_CUDA
32 #include <btas/array_adaptor.h>
33 #include <btas/varray/varray.h>
41 using cuda_um_thrust_vector =
42 thrust::device_vector<T, TiledArray::cuda_um_allocator<T>>;
45 template <MemorySpace Space,
typename Storage>
54 template <ExecutionSpace Space,
typename Storage>
57 case ExecutionSpace::CPU: {
60 using value_type =
typename Storage::value_type;
61 if (cudaEnv::instance()->concurrent_managed_access()) {
62 CudaSafeCall(cudaMemPrefetchAsync(data(vec),
63 size(vec) *
sizeof(value_type),
64 cudaCpuDeviceId, stream));
71 using value_type =
typename Storage::value_type;
73 if (cudaEnv::instance()->concurrent_managed_access()) {
74 CudaSafeCall(cudaGetDevice(&device));
75 CudaSafeCall(cudaMemPrefetchAsync(
76 data(vec), size(vec) *
sizeof(value_type), device, stream));
81 throw std::runtime_error(
"invalid execution space");
92 template <
typename Storage>
94 const cudaStream_t& stream = 0) {
96 TiledArray::to_execution_space<TiledArray::ExecutionSpace::CUDA>(storage,
106 template <
typename Storage>
107 typename Storage::value_type*
device_data(Storage& storage) {
108 return storage.data();
117 template <
typename Storage>
118 const typename Storage::value_type*
device_data(
const Storage& storage) {
119 return storage.data();
128 template <
class Archive,
typename T>
129 struct ArchiveLoadImpl;
130 template <
class Archive,
typename T>
131 struct ArchiveStoreImpl;
133 template <
class Archive,
typename T>
134 struct ArchiveLoadImpl<Archive,
TiledArray::cuda_um_thrust_vector<T>> {
135 static inline void load(
const Archive& ar,
136 TiledArray::cuda_um_thrust_vector<T>& x) {
137 typename thrust::device_vector<
138 T, TiledArray::cuda_um_allocator<T>>::size_type n(0);
141 for (
auto& xi : x) ar& xi;
145 template <
class Archive,
typename T>
146 struct ArchiveStoreImpl<Archive,
TiledArray::cuda_um_thrust_vector<T>> {
147 static inline void store(
const Archive& ar,
148 const TiledArray::cuda_um_thrust_vector<T>& x) {
150 for (
const auto& xi : x) ar& xi;
154 template <
class Archive,
typename T>
155 struct ArchiveLoadImpl<Archive,
TiledArray::cuda_um_btas_varray<T>> {
156 static inline void load(
const Archive& ar,
157 TiledArray::cuda_um_btas_varray<T>& x) {
158 typename TiledArray::cuda_um_btas_varray<T>::size_type n(0);
161 for (
auto& xi : x) ar& xi;
165 template <
class Archive,
typename T>
166 struct ArchiveStoreImpl<Archive,
TiledArray::cuda_um_btas_varray<T>> {
167 static inline void store(
const Archive& ar,
168 const TiledArray::cuda_um_btas_varray<T>& x) {
170 for (
const auto& xi : x) ar& xi;
177 #endif // TILEDARRAY_HAS_CUDA
179 #endif // TILEDARRAY_CUDA_UM_VECTOR_H__INCLUDED