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