KokkosMemoryPool.h

Go to the documentation of this file.

//* This file is part of the MOOSE framework
//* https://www.mooseframework.org
//*
//* All rights reserved, see COPYRIGHT for full restrictions
//* https://github.com/idaholab/moose/blob/master/COPYRIGHT
//*
//* Licensed under LGPL 2.1, please see LICENSE for details
//* https://www.gnu.org/licenses/lgpl-2.1.html
 
#pragma once
 
#include "KokkosHeader.h"
#include "KokkosArray.h"
 
#include "MooseApp.h"
 
namespace Moose::Kokkos
{
 
template <typename T>
class MemoryChunk
{
public:
  KOKKOS_FUNCTION MemoryChunk(const ::Kokkos::MemoryPool<MemSpace> & pool, dof_id_type size)
    : _pool(pool), _size(size)
  {
    auto ptr = pool.allocate(size);
 
    KOKKOS_ASSERT(ptr);
 
    _ptr = static_cast<T *>(ptr);
  }
  KOKKOS_FUNCTION ~MemoryChunk() { _pool.deallocate(_ptr, _size); }
  KOKKOS_FUNCTION T * get() const { return _ptr; }
 
private:
  const ::Kokkos::MemoryPool<MemSpace> & _pool;
  const dof_id_type _size;
  T * _ptr;
};
 
class MemoryPool
{
public:
  MemoryPool(dof_id_type size, unsigned int ways);
 
  template <typename T>
  KOKKOS_FUNCTION MemoryChunk<T> allocate(dof_id_type idx, unsigned int size) const
  {
    auto pool = idx % _pools.size();
 
    return MemoryChunk<T>(_pools[pool], size * sizeof(T));
  }
 
private:
  Array<::Kokkos::MemoryPool<MemSpace>> _pools;
};
 
class MemoryPoolHolder
{
public:
  MemoryPoolHolder(const MooseApp & app) : _app(app), _pool(_app.getKokkosMemoryPool()) {}
  MemoryPoolHolder(const MemoryPoolHolder & holder)
    : _app(holder._app), _pool(_app.getKokkosMemoryPool())
  {
  }
 
protected:
  KOKKOS_FUNCTION const MemoryPool & kokkosMemoryPool() const { return _pool; }
 
private:
  const MooseApp & _app;
  MemoryPool _pool;
};
 
} // namespace Moose::Kokkos