Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

pool_memory_resource optimization: disable tracking allocated blocks by default #702

Merged
merged 1 commit into from
Feb 16, 2021
Merged

pool_memory_resource optimization: disable tracking allocated blocks by default #702

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 17 additions & 6 deletions include/rmm/mr/device/pool_memory_resource.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand All @@ -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{};
Expand All @@ -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());

Expand All @@ -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
}

/**
Expand All @@ -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;
}
Expand Down Expand Up @@ -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();
}
Expand Down Expand Up @@ -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
Expand Down