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

arena_memory_resource optimization: disable tracking allocated blocks by default #732

Merged
merged 1 commit into from
Mar 17, 2021
Merged

arena_memory_resource optimization: disable tracking allocated blocks by default #732

Show file tree
Hide file tree
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
6 changes: 6 additions & 0 deletions include/rmm/mr/device/arena_memory_resource.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,11 +148,16 @@ class arena_memory_resource final : public device_memory_resource {
if (p == nullptr || bytes <= 0) return;

bytes = detail::arena::align_up(bytes);
#ifdef RMM_POOL_TRACK_ALLOCATIONS
if (!get_arena(stream).deallocate(p, bytes, stream)) {
deallocate_from_other_arena(p, bytes, stream);
}
#else
get_arena(stream).deallocate(p, bytes, stream);
#endif
}

#ifdef RMM_POOL_TRACK_ALLOCATIONS
/**
* @brief Deallocate memory pointed to by `p` that was allocated in a different arena.
*
Expand Down Expand Up @@ -186,6 +191,7 @@ class arena_memory_resource final : public device_memory_resource {
// global arena.
global_arena_.deallocate({p, bytes});
}
#endif

/**
* @brief Get the arena associated with the current thread or the given stream.
Expand Down
14 changes: 14 additions & 0 deletions include/rmm/mr/device/detail/arena.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -446,7 +446,9 @@ class arena {
{
lock_guard lock(mtx_);
auto const b = get_block(bytes);
#ifdef RMM_POOL_TRACK_ALLOCATIONS
allocated_blocks_.emplace(b.pointer(), b);
#endif
return b.pointer();
}

Expand All @@ -462,14 +464,19 @@ class arena {
bool deallocate(void* p, std::size_t bytes, cuda_stream_view stream)
{
lock_guard lock(mtx_);
#ifdef RMM_POOL_TRACK_ALLOCATIONS
auto const b = free_block(p, bytes);
#else
block const b{p, bytes};
#endif
if (b.is_valid()) {
auto const merged = coalesce_block(free_blocks_, b);
shrink_arena(merged, stream);
}
return b.is_valid();
}

#ifdef RMM_POOL_TRACK_ALLOCATIONS
/**
* @brief Deallocate memory pointed to by `p`, keeping all free superblocks.
*
Expand All @@ -489,6 +496,7 @@ class arena {
if (b.is_valid()) { global_arena_.deallocate(b); }
return b.is_valid();
}
#endif

/**
* @brief Clean the arena and deallocate free blocks from the global arena.
Expand All @@ -500,7 +508,9 @@ class arena {
lock_guard lock(mtx_);
global_arena_.deallocate(free_blocks_);
free_blocks_.clear();
#ifdef RMM_POOL_TRACK_ALLOCATIONS
allocated_blocks_.clear();
#endif
}

private:
Expand Down Expand Up @@ -537,6 +547,7 @@ class arena {
return global_arena_.allocate(superblock_size);
}

#ifdef RMM_POOL_TRACK_ALLOCATIONS
/**
* @brief Finds, frees and returns the block associated with pointer `p`.
*
Expand All @@ -558,6 +569,7 @@ class arena {

return found;
}
#endif

/**
* @brief Shrink this arena by returning free superblocks to upstream.
Expand All @@ -580,8 +592,10 @@ class arena {
global_arena<Upstream>& global_arena_;
/// Free blocks.
std::set<block> free_blocks_;
#ifdef RMM_POOL_TRACK_ALLOCATIONS
//// Map of pointer address to allocated blocks.
std::unordered_map<void*, block> allocated_blocks_;
#endif
/// Mutex for exclusive lock.
mutable std::mutex mtx_;
};
Expand Down