https://mooseframework.inl.gov
KokkosMemoryPool.h
Go to the documentation of this file.
1 //* This file is part of the MOOSE framework
2 //* https://www.mooseframework.org
3 //*
4 //* All rights reserved, see COPYRIGHT for full restrictions
5 //* https://github.com/idaholab/moose/blob/master/COPYRIGHT
6 //*
7 //* Licensed under LGPL 2.1, please see LICENSE for details
8 //* https://www.gnu.org/licenses/lgpl-2.1.html
9 
10 #pragma once
11 
12 #include "KokkosHeader.h"
13 #include "KokkosArray.h"
14 
15 #include "MooseApp.h"
16 
17 namespace Moose
18 {
19 namespace Kokkos
20 {
21 
30 template <typename T>
32 {
33 public:
39  KOKKOS_FUNCTION MemoryChunk(const ::Kokkos::MemoryPool<MemSpace> & pool, dof_id_type size)
40  : _pool(pool), _size(size)
41  {
42  auto ptr = pool.allocate(size);
43 
44  KOKKOS_ASSERT(ptr);
45 
46  _ptr = static_cast<T *>(ptr);
47  }
51  KOKKOS_FUNCTION ~MemoryChunk() { _pool.deallocate(_ptr, _size); }
56  KOKKOS_FUNCTION T * get() const { return _ptr; }
57 
58 private:
62  const ::Kokkos::MemoryPool<MemSpace> & _pool;
70  T * _ptr;
71 };
72 
78 {
79 public:
85  MemoryPool(dof_id_type size, unsigned int ways);
86 
94  template <typename T>
95  KOKKOS_FUNCTION MemoryChunk<T> allocate(dof_id_type idx, unsigned int size) const
96  {
97  auto pool = idx % _pools.size();
98 
99  return MemoryChunk<T>(_pools[pool], size * sizeof(T));
100  }
101 
102 private:
107 };
108 
114 {
115 public:
120  MemoryPoolHolder(const MooseApp & app) : _app(app), _pool(_app.getKokkosMemoryPool()) {}
125  : _app(holder._app), _pool(_app.getKokkosMemoryPool())
126  {
127  }
128 
129 protected:
134  KOKKOS_FUNCTION const MemoryPool & kokkosMemoryPool() const { return _pool; }
135 
136 private:
140  const MooseApp & _app;
145 };
146 
147 } // namespace Kokkos
148 } // namespace Moose
const MooseApp & _app
Reference of the MOOSE app.
The Kokkos array class.
Definition: KokkosArray.h:56
T * _ptr
Pointer to the memory chunk.
A temporary object returned by the Kokkos memory pool during memory chunk allocation.
KOKKOS_FUNCTION MemoryChunk(const ::Kokkos::MemoryPool< MemSpace > &pool, dof_id_type size)
Constructor.
KOKKOS_FUNCTION const MemoryPool & kokkosMemoryPool() const
Get the const reference of the Kokkos memory pool.
The Kokkos interface that holds the Kokkos memory pool.
Base class for MOOSE-based applications.
Definition: MooseApp.h:103
MemoryPoolHolder(const MemoryPoolHolder &holder)
Copy constructor.
The Kokkos class that manages memory pool for dynamically-sized temporary arrays in Kokkos parallel f...
MemoryPool _pool
Device copy of the Kokkos memory pool.
Array<::Kokkos::MemoryPool< MemSpace > > _pools
Kokkos memory pools.
KOKKOS_FUNCTION MemoryChunk< T > allocate(dof_id_type idx, unsigned int size) const
Allocate a memory chunk The returned object should be alive while the memory chunk is used...
const ::Kokkos::MemoryPool< MemSpace > & _pool
Reference of the Kokkos memory pool.
MOOSE now contains C++17 code, so give a reasonable error message stating what the user can do to add...
KOKKOS_FUNCTION ~MemoryChunk()
Destructor.
MemoryPool(dof_id_type size, unsigned int ways)
Constructor.
MemoryPoolHolder(const MooseApp &app)
Constructor.
unsigned int idx(const ElemType type, const unsigned int nx, const unsigned int i, const unsigned int j)
uint8_t dof_id_type
const dof_id_type _size
Memory chunk size in the number of bytes.