memory_pool.hpp

// FEAT3: Finite Element Analysis Toolbox, Version 3
// Copyright (C) 2010 by Stefan Turek & the FEAT group
// FEAT3 is released under the GNU General Public License version 3,
// see the file 'copyright.txt' in the top level directory for details.
 
#pragma once
 
// includes, FEAT
#include <kernel/base_header.hpp>
#include <kernel/util/exception.hpp>
#include <kernel/util/assertion.hpp>
#include <kernel/util/cuda_util.hpp>
#include <kernel/util/random.hpp>
#include <kernel/backend.hpp>
 
#include <map>
#include <cstring>
#include <typeinfo>
#include <cstdio>
#include <cstddef>
 
#ifdef FEAT_HAVE_MKL
FEAT_DISABLE_WARNINGS
#include <mkl.h>
FEAT_RESTORE_WARNINGS
#endif
 
 
namespace FEAT
{
  namespace Util
  {
    namespace Intern
    {
      struct MemoryInfo
      {
        Index counter;
        Index size;
      };
    }
  } // namespace Util
 
    class MemoryPool
    {
      private:
        static std::map<void*, Util::Intern::MemoryInfo> _pool;
 
      public:
 
        static void initialize()
        {
        }
 
        static void finalize()
        {
          if (_pool.size() > 0)
          {
            std::cout << stderr << " Error: MemoryPool still contains memory chunks on deconstructor call\n";
            std::exit(1);
          }
 
#ifdef FEAT_HAVE_MKL
          mkl_free_buffers();
#endif
        }
 
        template <typename DT_>
        static DT_ * allocate_memory(Index count)
        {
          DT_ * memory(nullptr);
 
          if (count == 0)
            return memory;
 
          if (count%4 != 0)
            count = count + (4ul - count%4);
 
#ifdef FEAT_HAVE_CUDA
          memory = (DT_*)Util::cuda_malloc_managed(count * sizeof(DT_));
#else
          memory = (DT_*)::malloc(count * sizeof(DT_));
#endif
          if (memory == nullptr)
            XABORTM("MemoryPool allocation error!");
 
          Util::Intern::MemoryInfo mi;
          mi.counter = 1;
          mi.size = count * sizeof(DT_);
          _pool.insert(std::pair<void*, Util::Intern::MemoryInfo>(memory, mi));
 
          return memory;
        }
 
        static void increase_memory(void * address)
        {
          XASSERT(address != nullptr);
 
          std::map<void*, Util::Intern::MemoryInfo>::iterator it(_pool.find(address));
          if (it != _pool.end())
          {
            it->second.counter = it->second.counter + 1;
            return;
          }
 
          XABORTM("MemoryPool::increase_memory: Memory address not found!");
        }
 
        static void release_memory(void * address)
        {
          if (address == nullptr)
            return;
 
          std::map<void*, Util::Intern::MemoryInfo>::iterator it(_pool.find(address));
          if (it != _pool.end())
          {
            if(it->second.counter == 1)
            {
#ifdef FEAT_HAVE_CUDA
              Util::cuda_free(address);
#else
              ::free(address);
#endif
              _pool.erase(it);
            }
            else
            {
              it->second.counter = it->second.counter - 1;
            }
            return;
          }
 
          XABORTM("MemoryPool::release_memory: Memory address not found!");
        }
 
        template <typename DT_>
        [[deprecated("no download necessary in unified memory environment.")]]
        inline static void download(DT_ * dest, const DT_ * const src, const Index count)
        {
          if (dest == src)
            return;
 
          ::memcpy(dest, src, count * sizeof(DT_));
        }
 
        template <typename DT_>
        [[deprecated("no upload necessary in unified memory environment.")]]
        inline static void upload(DT_ * dest, const DT_ * const src, const Index count)
        {
          if (dest == src)
            return;
 
          ::memcpy(dest, src, count * sizeof(DT_));
        }
 
        template <typename DT_>
        [[deprecated("no get_element necessary in unified memory environment.")]]
        inline static const DT_ & get_element(const DT_ * data, const Index index)
        {
          return data[index];
        }
 
        template <typename DT_>
        static void set_memory(DT_ * address, const DT_ val, const Index count = 1)
        {
          switch(Backend::get_preferred_backend())
          {
#ifdef FEAT_HAVE_CUDA
            case PreferredBackend::cuda:
            {
              FEAT::Util::cuda_set_memory(address, val, count);
              return;
            }
#endif
            case PreferredBackend::generic:
            default:
            {
              FEAT_PRAGMA_OMP(parallel for)
              for (Index i = 0 ; i < count ; ++i)
              {
                address[i] = val;
              }
              return;
            }
          }
        }
 
        template <typename DT_>
        static void set_memory(Random& rng, DT_ min, DT_ max, DT_ * address, const Index count = 1)
        {
          switch(Backend::get_preferred_backend())
          {
#ifdef FEAT_HAVE_CUDA
            case PreferredBackend::cuda:
            {
              std::vector<DT_> tmp(count);
              std::generate(tmp.begin(), tmp.end(), [&](){return rng(min,max);});
              FEAT::Util::cuda_copy(address, tmp.data(), count*sizeof(DT_));
              return;
            }
#endif
            case PreferredBackend::generic:
            default:
            {
              // we don't use OpenMP here, because this would result in a non-deterministic vector
              for (Index i(0) ; i < count ; ++i)
              {
                address[i] = rng(min, max);
              }
              return;
            }
          }
        }
 
        template <typename DT_>
        static void copy(DT_ * dest, const DT_ * src, const Index count)
        {
          if (dest == src)
            return;
 
          switch(Backend::get_preferred_backend())
          {
#ifdef FEAT_HAVE_CUDA
            case PreferredBackend::cuda:
            {
              FEAT::Util::cuda_copy(dest, src, count * sizeof(DT_));
              return;
            }
#endif
            case PreferredBackend::generic:
            default:
            {
#ifdef FEAT_HAVE_OMP
              FEAT_PRAGMA_OMP(parallel for)
              for (Index i = 0 ; i < count ; ++i)
              {
                dest[i] = src[i];
              }
#else
              ::memcpy(dest, src, count * sizeof(DT_));
#endif
              return;
            }
          }
        }
 
        template <typename DT_>
        static void convert(DT_ * dest, const DT_ * src, const Index count)
        {
          if (dest == src)
            return;
 
#ifdef FEAT_HAVE_OMP
          FEAT_PRAGMA_OMP(parallel for)
          for (Index i = 0 ; i < count ; ++i)
          {
            dest[i] = src[i];
          }
#else
          ::memcpy(dest, src, count * sizeof(DT_));
#endif
        }
 
        template <typename DT1_, typename DT2_>
        static void convert(DT1_ * dest, const DT2_ * src, const Index count)
        {
 
          switch(Backend::get_preferred_backend())
          {
#ifdef FEAT_HAVE_CUDA
            case PreferredBackend::cuda:
            {
              FEAT::Util::cuda_convert(dest, src, count);
              return;
            }
#endif
            case PreferredBackend::generic:
            default:
            {
              FEAT_PRAGMA_OMP(parallel for)
              for (Index i = 0 ; i < count ; ++i)
              {
                dest[i] = DT1_(src[i]);
              }
              return;
            }
          }
        }
 
        NOINLINE static void synchronize()
        {
#ifdef FEAT_HAVE_CUDA
          //avoid calling expensive synchronize if cuda is not actively used
          if (Backend::get_preferred_backend() == PreferredBackend::cuda)
            FEAT::Util::cuda_synchronize();
#endif
        }
 
        static Index allocated_memory()
        {
          Index bytes(0);
          for (auto& i : _pool)
          {
            bytes += i.second.size;
          }
          return bytes;
        }
 
        static Index allocated_size(void * address)
        {
          std::map<void*, Util::Intern::MemoryInfo>::iterator it(_pool.find(address));
          if (it != _pool.end())
          {
            return it->second.size;
          }
          else
            XABORTM("MemoryPool::allocated_size: Memory address not found!");
        }
    };
 
} // namespace FEAT