XGC1
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Pages
loadable_buffer.hpp
Go to the documentation of this file.
1 #ifndef LOADABLE_BUFFER_HPP
2 #define LOADABLE_BUFFER_HPP
3 
4 #include "space_settings.hpp"
5 
6 /* This class facilitates packing Kokkos Views of arbitrary dimension into a single contiguous allocation,
7  * which can be useful for data movement such as MPI operations. The buffer must be unloaded in the reverse order
8  * in which it is loaded (first in, last out).
9  * Currently requires that the incoming Views are of type T.
10  */
11 template<typename T, class Device>
13  View<T*, CLayout, Device> buffer;
14  int offset;
15 
16  /* size_sum is used to determine the necessary buffer size when using the constructor that takes a template pack of Views
17  * The final template has only one View left, so returns the size of that View.
18  */
19  template<class V>
20  int size_sum(V first) {
21  return first.size();
22  }
23 
24  /* size_sum is used to determine the necessary buffer size when using the constructor that takes a template pack of Views
25  * The generic template adds the first View in the template pack, then returns the size_sum of the remaining Views in the pack
26  */
27  template<class V, class... TRest>
28  int size_sum(V first, TRest... args) {
29  return first.size() + size_sum(args...);
30  }
31 
32  public:
33 
34  /* Constructor where the user specifies the size manually.
35  */
37  : buffer("buffer",size),
38  offset(0)
39  {}
40 
41  /* Constructor where the user inputs a list of Views, the size of which is added to determine how big the buffer needs to be.
42  * This is essentially 'previewing' the loading operations to get the allocation size required.
43  */
44  template<class... Vs>
45  LoadableBuffer(Vs... views)
46  : buffer("buffer",size_sum(views...)),
47  offset(0)
48  {}
49 
50  /* Loads the input View into the buffer, and offsets the buffer to prepare for the next View. If do_copy is false,
51  * then only the offset is performed. This saves an unnecessary copy if the copied value would have been unused, e.g.
52  * for a broadcast, but is necessary so that the unload can be done properly.
53  */
54  template<class V>
55  void load(V view, bool do_copy){
56  int size_of_view = view.size();
57  View<T*, CLayout, Device, Kokkos::MemoryTraits<Kokkos::Unmanaged>> offset_buffer(buffer.data() + offset, size_of_view);
58  View<T*, CLayout, typename V::device_type, Kokkos::MemoryTraits<Kokkos::Unmanaged>> unmanaged_view(view.data(), size_of_view);
59 
60  if(do_copy) Kokkos::deep_copy(offset_buffer, unmanaged_view);
61  offset += size_of_view;
62  }
63 
64  /* Unloads the buffer into the input View, and offsets the buffer to prepare for the next View. If do_copy is false,
65  * then only the offset is performed. This saves an unnecessary copy if the copied value would have been unused, e.g.
66  * for a broadcast.
67  */
68  template<class V>
69  void unload(V view, bool do_copy){
70  int size_of_view = view.size();
71  offset -= size_of_view;
72  View<T*, CLayout, Device, Kokkos::MemoryTraits<Kokkos::Unmanaged>> offset_buffer(buffer.data() + offset, size_of_view);
73  View<T*, CLayout, typename V::device_type, Kokkos::MemoryTraits<Kokkos::Unmanaged>> unmanaged_view(view.data(), size_of_view);
74 
75  if(do_copy) Kokkos::deep_copy(unmanaged_view, offset_buffer);
76  }
77 
78  /* Returns the size of the buffer.
79  */
80  int size(){
81  return buffer.size();
82  }
83 
84  /* Returns the pointer to the buffer
85  */
86  T* data(){
87  return buffer.data();
88  }
89 };
90 
91 #endif
LoadableBuffer(Vs...views)
Definition: loadable_buffer.hpp:45
View< T *, CLayout, Device > buffer
The buffer allocation.
Definition: loadable_buffer.hpp:13
void load(V view, bool do_copy)
Definition: loadable_buffer.hpp:55
int offset
Tracks where the most recently loaded View ended in the buffer.
Definition: loadable_buffer.hpp:14
int size_sum(V first)
Definition: loadable_buffer.hpp:20
Definition: loadable_buffer.hpp:12
void unload(V view, bool do_copy)
Definition: loadable_buffer.hpp:69
T * data()
Definition: loadable_buffer.hpp:86
int size_sum(V first, TRest...args)
Definition: loadable_buffer.hpp:28
int size()
Definition: loadable_buffer.hpp:80
LoadableBuffer(int size)
Definition: loadable_buffer.hpp:36