Skip to content

Commit 8af44bd

Browse files
authored
Merge pull request #472 from ValeevGroup/evaleev/fix/umpire-allocator-page-size
device::Env::initialize: use correct page sizes for Umpire allocators
2 parents 236af70 + 086e1e4 commit 8af44bd

File tree

1 file changed

+10
-7
lines changed

1 file changed

+10
-7
lines changed

src/TiledArray/external/device.h

+10-7
Original file line numberDiff line numberDiff line change
@@ -503,8 +503,7 @@ class Env {
503503
/// \param page_size memory added to the pools supporting `this->um_allocator()`, `this->device_allocator()`, and `this->pinned_allocator()` in chunks of at least
504504
/// this size (bytes) [default=2^25]
505505
/// \param pinned_alloc_limit the maximum total amount of memory (in bytes) that
506-
/// allocator returned by `this->pinned_allocator()` can allocate;
507-
/// this allocator is not used by default [default=0]
506+
/// allocator returned by `this->pinned_allocator()` can allocate [default=2^40]
508507
// clang-format on
509508
static void initialize(World& world = TiledArray::get_default_world(),
510509
const std::uint64_t page_size = (1ul << 25),
@@ -563,8 +562,9 @@ class Env {
563562
// allocate all currently-free memory for UM pool
564563
auto um_dynamic_pool =
565564
rm.makeAllocator<umpire::strategy::QuickPool, introspect>(
566-
"UMDynamicPool", rm.getAllocator("UM"), mem_total_free.second,
567-
pinned_alloc_limit);
565+
"UMDynamicPool", rm.getAllocator("UM"),
566+
/* first_minimum_pool_allocation_size = */ 0,
567+
/* next_minimum_pool_allocation_size = */ page_size);
568568

569569
// allocate zero memory for device pool
570570
auto dev_size_limited_alloc =
@@ -573,8 +573,9 @@ class Env {
573573
mem_total_free.first);
574574
auto dev_dynamic_pool =
575575
rm.makeAllocator<umpire::strategy::QuickPool, introspect>(
576-
"DEVICEDynamicPool", dev_size_limited_alloc, 0,
577-
pinned_alloc_limit);
576+
"DEVICEDynamicPool", dev_size_limited_alloc,
577+
/* first_minimum_pool_allocation_size = */ 0,
578+
/* next_minimum_pool_allocation_size = */ page_size);
578579

579580
// allocate pinned_alloc_limit in pinned memory
580581
auto pinned_size_limited_alloc =
@@ -584,7 +585,9 @@ class Env {
584585
auto pinned_dynamic_pool =
585586
rm.makeAllocator<umpire::strategy::QuickPool, introspect>(
586587
"QuickPool_SizeLimited_PINNED", pinned_size_limited_alloc,
587-
page_size, page_size, /* alignment */ TILEDARRAY_ALIGN_SIZE);
588+
/* first_minimum_pool_allocation_size = */ 0,
589+
/* next_minimum_pool_allocation_size = */ page_size,
590+
/* alignment */ TILEDARRAY_ALIGN_SIZE);
588591

589592
auto env = std::unique_ptr<Env>(new Env(
590593
world, num_visible_devices, compute_devices, num_streams_per_device,

0 commit comments

Comments
 (0)