Line data    Source code

       1             : #ifndef AMREX_ARENA_H_
       2             : #define AMREX_ARENA_H_
       3             : #include <AMReX_Config.H>
       4             : 
       5             : #include <AMReX_BLassert.H>
       6             : #include <AMReX_INT.H>
       7             : #include <cstddef>
       8             : #include <cstdlib>
       9             : #include <limits>
      10             : #include <utility>
      11             : 
      12             : namespace amrex {
      13             : 
      14             : /**
      15             :  * \brief Given a minimum required size of size bytes, this returns
      16             :  * the next largest arena size that will align to align_requirement bytes
      17             :  */
      18           0 : inline std::size_t aligned_size (std::size_t align_requirement, std::size_t size) noexcept
      19             : {
      20           0 :     return ((size + (align_requirement-1)) / align_requirement) * align_requirement;
      21             : }
      22             : 
      23             : inline bool is_aligned (const void* p, std::size_t alignment) noexcept
      24             : {
      25             :     auto* q = const_cast<void*>(p);
      26             :     auto space = alignment;
      27             :     return std::align(alignment, alignment, q, space) == p;
      28             : }
      29             : 
      30             : class Arena;
      31             : 
      32             : Arena* The_Arena ();
      33             : Arena* The_Async_Arena ();
      34             : Arena* The_Device_Arena ();
      35             : Arena* The_Managed_Arena ();
      36             : Arena* The_Pinned_Arena ();
      37             : Arena* The_Comms_Arena ();
      38             : Arena* The_Cpu_Arena ();
      39             : 
      40             : struct ArenaInfo
      41             : {
      42             :     Long release_threshold = std::numeric_limits<Long>::max();
      43             :     bool use_cpu_memory = false;
      44             :     bool device_use_managed_memory = true;
      45             :     bool device_set_readonly = false;
      46             :     bool device_set_preferred = false;
      47             :     bool device_use_hostalloc = false;
      48             :     ArenaInfo& SetReleaseThreshold (Long rt) noexcept {
      49             :         release_threshold = rt;
      50             :         return *this;
      51             :     }
      52             :     ArenaInfo& SetDeviceMemory () noexcept {
      53             :         device_use_managed_memory = false;
      54             :         device_use_hostalloc = false;
      55             :         return *this;
      56             :     }
      57             :     ArenaInfo& SetReadOnly () noexcept {
      58             :         BL_ASSERT(device_use_managed_memory);
      59             :         device_set_readonly = true;
      60             :         return *this;
      61             :     }
      62             :     ArenaInfo& SetPreferred () noexcept {
      63             :         BL_ASSERT(device_use_managed_memory);
      64             :         device_set_preferred = true;
      65             :         return *this;
      66             :     }
      67             :     ArenaInfo& SetHostAlloc () noexcept {
      68             :         device_use_hostalloc = true;
      69             :         device_use_managed_memory = false;
      70             :         return *this;
      71             :     }
      72             :     ArenaInfo& SetCpuMemory () noexcept {
      73             :         use_cpu_memory = true;
      74             :         device_use_managed_memory = false;
      75             :         device_set_readonly = false;
      76             :         device_set_preferred = false;
      77             :         device_use_hostalloc = false;
      78             :         return *this;
      79             :     }
      80             : };
      81             : 
      82             : /**
      83             : * \brief
      84             : * A virtual base class for objects that manage their own dynamic
      85             : * memory allocation.
      86             : */
      87             : class Arena
      88             : {
      89             : public:
      90             : 
      91             :     virtual ~Arena () = default;
      92             :     Arena () noexcept = default;
      93             : 
      94             :     Arena (const Arena& rhs) = delete;
      95             :     Arena (Arena&& rhs) = delete;
      96             :     Arena& operator= (const Arena& rhs) = delete;
      97             :     Arena& operator= (Arena&& rhs) = delete;
      98             : 
      99             :     /**
     100             :     * Allocate a dynamic memory arena
     101             :     * \param sz size of the memory request
     102             :     * \return a pointer to the allocated memory
     103             :     */
     104             :     [[nodiscard]] virtual void* alloc (std::size_t sz) = 0;
     105             : 
     106             :     /**
     107             :      * Try to allocate in-place by extending the capacity of given pointer.
     108             :      */
     109             :     [[nodiscard]] virtual std::pair<void*,std::size_t>
     110             :     alloc_in_place (void* /*pt*/, std::size_t /*szmin*/, std::size_t szmax)
     111             :     {
     112             :         auto* p = alloc(szmax);
     113             :         return std::make_pair(p, szmax);
     114             :     }
     115             : 
     116             :     /**
     117             :      * Try to shrink in-place
     118             :      */
     119             :     [[nodiscard]] virtual void*
     120             :     shrink_in_place (void* /*pt*/, std::size_t sz)
     121             :     {
     122             :         return alloc(sz);
     123             :     }
     124             : 
     125             :     /**
     126             :     * \brief A pure virtual function for deleting the arena pointed to by pt
     127             :     */
     128             :     virtual void free (void* pt) = 0;
     129             : 
     130             :     /**
     131             :     * \brief Free unused memory back to the system.  Return value is the
     132             :     * amount memory freed.
     133             :     */
     134             :     virtual std::size_t freeUnused () { return 0; }
     135             : 
     136             :     // isDeviceAccessible and isHostAccessible can both be true.
     137             :     [[nodiscard]] virtual bool isDeviceAccessible () const;
     138             :     [[nodiscard]] virtual bool isHostAccessible () const;
     139             : 
     140             :     // Note that isManaged, isDevice and isPinned are mutually exclusive.
     141             :     // For memory allocated by cudaMalloc* etc., one of them returns true.
     142             :     // Otherwise, neither is true.
     143             :     [[nodiscard]] virtual bool isManaged () const;
     144             :     [[nodiscard]] virtual bool isDevice () const;
     145             :     [[nodiscard]] virtual bool isPinned () const;
     146             : 
     147             :     /**
     148             :      * \brief Does the device have enough free memory for allocating this
     149             :      * much memory?  For CPU builds, this always return true.  This is not a
     150             :      * const function because it may attempt to release memory back to the
     151             :      * system.
     152             :      */
     153             :     [[nodiscard]] virtual bool hasFreeDeviceMemory (std::size_t sz);
     154             : 
     155             :     /**
     156             :      * \brief Add this Arena to the list of Arenas that are profiled by TinyProfiler.
     157             :      * \param memory_name The name of this arena in the TinyProfiler output.
     158             :      */
     159             :     virtual void registerForProfiling (const std::string& memory_name);
     160             : 
     161             : #ifdef AMREX_USE_GPU
     162             :     //! Is this GPU stream ordered memory allocator?
     163             :     [[nodiscard]] virtual bool isStreamOrderedArena () const { return false; }
     164             : #endif
     165             : 
     166             :     /**
     167             :     * \brief Given a minimum required arena size of sz bytes, this returns
     168             :     * the next largest arena size that will align to align_size bytes
     169             :     */
     170             :     static std::size_t align (std::size_t sz);
     171             : 
     172             :     static void Initialize ();
     173             :     static void PrintUsage ();
     174             :     static void PrintUsageToFiles (std::string const& filename, std::string const& message);
     175             :     static void Finalize ();
     176             : 
     177             : #if 0
     178             :     union Word
     179             :     {
     180             :         void*  p;
     181             :         long long ll;
     182             :         long double ld;
     183             :         void (*f) ();
     184             :     };
     185             :     static const std::size_t align_size = sizeof(Word);
     186             : #endif
     187             : 
     188             :     static const std::size_t align_size = 16;
     189             : 
     190             :     /**
     191             :      *  \brief Return the ArenaInfo object for querying
     192             :      */
     193             :     [[nodiscard]] const ArenaInfo& arenaInfo () const { return arena_info; }
     194             : 
     195             : protected:
     196             : 
     197             :     ArenaInfo arena_info;
     198             : 
     199             :     virtual std::size_t freeUnused_protected () { return 0; }
     200             :     void* allocate_system (std::size_t nbytes);
     201             :     void deallocate_system (void* p, std::size_t nbytes);
     202             : };
     203             : 
     204             : }
     205             : 
     206             : #endif /*BL_ARENA_H*/