From 77429624dea7afd683a4d940647a9ccb24063c64 Mon Sep 17 00:00:00 2001
From: Mark Harris <mharris@nvidia.com>
Date: Tue, 16 Feb 2021 09:47:24 +1100
Subject: [PATCH] Tracking allocated blocks in pool disabled by default

---
 .../rmm/mr/device/pool_memory_resource.hpp    | 23 ++++++++++++++-----
 1 file changed, 17 insertions(+), 6 deletions(-)

diff --git a/include/rmm/mr/device/pool_memory_resource.hpp b/include/rmm/mr/device/pool_memory_resource.hpp
index 4aa252528..038d7f02c 100644
--- a/include/rmm/mr/device/pool_memory_resource.hpp
+++ b/include/rmm/mr/device/pool_memory_resource.hpp
@@ -283,9 +283,8 @@ class pool_memory_resource final
 
     try {
       void* p = upstream_mr_->allocate(size, stream);
-      block_type b{reinterpret_cast<char*>(p), size, true};
-      upstream_blocks_.emplace_back(b);  // TODO: with C++17 use version that returns a reference
-      return thrust::optional<block_type>{b};
+      return thrust::optional<block_type>{
+        *upstream_blocks_.emplace(reinterpret_cast<char*>(p), size, true).first};
     } catch (std::exception const& e) {
       return thrust::nullopt;
     }
@@ -305,7 +304,9 @@ class pool_memory_resource final
   split_block allocate_from_block(block_type const& b, size_t size)
   {
     block_type const alloc{b.pointer(), size, b.is_head()};
+#ifdef RMM_POOL_TRACK_ALLOCATIONS
     allocated_blocks_.insert(alloc);
+#endif
 
     auto rest =
       (b.size() > size) ? block_type{b.pointer() + size, b.size() - size, false} : block_type{};
@@ -323,8 +324,8 @@ class pool_memory_resource final
    */
   block_type free_block(void* p, size_t size) noexcept
   {
+#ifdef RMM_POOL_TRACK_ALLOCATIONS
     if (p == nullptr) return block_type{};
-
     auto const i = allocated_blocks_.find(static_cast<char*>(p));
     RMM_LOGGING_ASSERT(i != allocated_blocks_.end());
 
@@ -333,6 +334,10 @@ class pool_memory_resource final
     allocated_blocks_.erase(i);
 
     return block;
+#else
+    auto const i = upstream_blocks_.find(static_cast<char*>(p));
+    return block_type{static_cast<char*>(p), size, (i != upstream_blocks_.end())};
+#endif
   }
 
   /**
@@ -355,7 +360,9 @@ class pool_memory_resource final
     for (auto b : upstream_blocks_)
       upstream_mr_->deallocate(b.pointer(), b.size());
     upstream_blocks_.clear();
+#ifdef RMM_POOL_TRACK_ALLOCATIONS
     allocated_blocks_.clear();
+#endif
 
     current_pool_size_ = 0;
   }
@@ -383,9 +390,11 @@ class pool_memory_resource final
     }
     std::cout << "total upstream: " << upstream_total << " B\n";
 
+#ifdef RMM_POOL_TRACK_ALLOCATIONS
     std::cout << "allocated_blocks: " << allocated_blocks_.size() << "\n";
     for (auto b : allocated_blocks_)
       b.print();
+#endif
 
     this->print_free_blocks();
   }
@@ -429,10 +438,12 @@ class pool_memory_resource final
   std::size_t current_pool_size_{};
   thrust::optional<std::size_t> maximum_pool_size_{};
 
+#ifdef RMM_POOL_TRACK_ALLOCATIONS
   std::set<block_type, rmm::mr::detail::compare_blocks<block_type>> allocated_blocks_;
+#endif
 
-  // blocks allocated from upstream: so they can be easily freed
-  std::vector<block_type> upstream_blocks_;
+  // blocks allocated from upstream
+  std::set<block_type, rmm::mr::detail::compare_blocks<block_type>> upstream_blocks_;
 };  // namespace mr
 
 }  // namespace mr