device::Env::initialize: use correct page sizes for Umpire allocation…

…s + do not allocate anything at the start
ValeevGroup · Sep 4, 2024 · 086e1e4 · 086e1e4
1 parent 236af70
commit 086e1e4
Showing 1 changed file with 10 additions and 7 deletions.
diff --git a/src/TiledArray/external/device.h b/src/TiledArray/external/device.h
@@ -503,8 +503,7 @@ class Env {
   /// \param page_size memory added to the pools supporting `this->um_allocator()`, `this->device_allocator()`, and `this->pinned_allocator()` in chunks of at least
   ///                  this size (bytes) [default=2^25]
   /// \param pinned_alloc_limit the maximum total amount of memory (in bytes) that
-  ///        allocator returned by `this->pinned_allocator()` can allocate;
-  ///        this allocator is not used by default [default=0]
+  ///        allocator returned by `this->pinned_allocator()` can allocate [default=2^40]
   // clang-format on
   static void initialize(World& world = TiledArray::get_default_world(),
                          const std::uint64_t page_size = (1ul << 25),
@@ -563,8 +562,9 @@ class Env {
       // allocate all currently-free memory for UM pool
       auto um_dynamic_pool =
           rm.makeAllocator<umpire::strategy::QuickPool, introspect>(
-              "UMDynamicPool", rm.getAllocator("UM"), mem_total_free.second,
-              pinned_alloc_limit);
+              "UMDynamicPool", rm.getAllocator("UM"),
+              /* first_minimum_pool_allocation_size = */ 0,
+              /* next_minimum_pool_allocation_size = */ page_size);
 
       // allocate zero memory for device pool
       auto dev_size_limited_alloc =
@@ -573,8 +573,9 @@ class Env {
               mem_total_free.first);
       auto dev_dynamic_pool =
           rm.makeAllocator<umpire::strategy::QuickPool, introspect>(
-              "DEVICEDynamicPool", dev_size_limited_alloc, 0,
-              pinned_alloc_limit);
+              "DEVICEDynamicPool", dev_size_limited_alloc,
+              /* first_minimum_pool_allocation_size = */ 0,
+              /* next_minimum_pool_allocation_size = */ page_size);
 
       // allocate pinned_alloc_limit in pinned memory
       auto pinned_size_limited_alloc =
@@ -584,7 +585,9 @@ class Env {
       auto pinned_dynamic_pool =
           rm.makeAllocator<umpire::strategy::QuickPool, introspect>(
               "QuickPool_SizeLimited_PINNED", pinned_size_limited_alloc,
-              page_size, page_size, /* alignment */ TILEDARRAY_ALIGN_SIZE);
+              /* first_minimum_pool_allocation_size = */ 0,
+              /* next_minimum_pool_allocation_size = */ page_size,
+              /* alignment */ TILEDARRAY_ALIGN_SIZE);
 
       auto env = std::unique_ptr<Env>(new Env(
           world, num_visible_devices, compute_devices, num_streams_per_device,