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

fix race condition in limiting resource adapter #869

Merged
merged 2 commits into from
Sep 14, 2021
Merged

fix race condition in limiting resource adapter #869

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
04529ef
fix race condition in limiting resource adapter
rongou Sep 13, 2021
72dbb4d
review feedback
rongou Sep 13, 2021
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
55 changes: 29 additions & 26 deletions include/rmm/mr/device/limiting_resource_adaptor.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,8 +21,7 @@

#include <cstddef>

namespace rmm {
namespace mr {
namespace rmm::mr {
/**
* @brief Resource that uses `Upstream` to allocate memory and limits the total
* allocations possible.
Expand Down Expand Up @@ -59,12 +58,12 @@ class limiting_resource_adaptor final : public device_memory_resource {
RMM_EXPECTS(nullptr != upstream, "Unexpected null upstream resource pointer.");
}

limiting_resource_adaptor() = delete;
~limiting_resource_adaptor() = default;
limiting_resource_adaptor(limiting_resource_adaptor const&) = delete;
limiting_resource_adaptor(limiting_resource_adaptor&&) = default;
limiting_resource_adaptor() = delete;
~limiting_resource_adaptor() override = default;
limiting_resource_adaptor(limiting_resource_adaptor const&) = delete;
limiting_resource_adaptor(limiting_resource_adaptor&&) noexcept = default;
limiting_resource_adaptor& operator=(limiting_resource_adaptor const&) = delete;
limiting_resource_adaptor& operator=(limiting_resource_adaptor&&) = default;
limiting_resource_adaptor& operator=(limiting_resource_adaptor&&) noexcept = default;

/**
* @brief Return pointer to the upstream resource.
Expand All @@ -79,14 +78,17 @@ class limiting_resource_adaptor final : public device_memory_resource {
* @return true The upstream resource supports streams
* @return false The upstream resource does not support streams.
*/
bool supports_streams() const noexcept override { return upstream_->supports_streams(); }
[[nodiscard]] bool supports_streams() const noexcept override
{
return upstream_->supports_streams();
}

/**
* @brief Query whether the resource supports the get_mem_info API.
*
* @return bool true if the upstream resource supports get_mem_info, false otherwise.
*/
bool supports_get_mem_info() const noexcept override
[[nodiscard]] bool supports_get_mem_info() const noexcept override
{
return upstream_->supports_get_mem_info();
}
Expand All @@ -100,7 +102,7 @@ class limiting_resource_adaptor final : public device_memory_resource {
* @return std::size_t number of bytes that have been allocated through this
* allocator.
*/
std::size_t get_allocated_bytes() const { return allocated_bytes_; }
[[nodiscard]] std::size_t get_allocated_bytes() const { return allocated_bytes_; }

/**
* @brief Query the maximum number of bytes that this allocator is allowed
Expand All @@ -109,7 +111,7 @@ class limiting_resource_adaptor final : public device_memory_resource {
*
* @return std::size_t max number of bytes allowed for this allocator
*/
std::size_t get_allocation_limit() const { return allocation_limit_; }
[[nodiscard]] std::size_t get_allocation_limit() const { return allocation_limit_; }

private:
/**
Expand All @@ -127,17 +129,19 @@ class limiting_resource_adaptor final : public device_memory_resource {
*/
void* do_allocate(std::size_t bytes, cuda_stream_view stream) override
{
void* p = nullptr;

std::size_t proposed_size = rmm::detail::align_up(bytes, allocation_alignment_);
if (proposed_size + allocated_bytes_ <= allocation_limit_) {
p = upstream_->allocate(bytes, stream);
allocated_bytes_ += proposed_size;
} else {
throw rmm::bad_alloc{"Exceeded memory limit"};
auto const proposed_size = rmm::detail::align_up(bytes, allocation_alignment_);
auto const old = allocated_bytes_.fetch_add(proposed_size);
if (old + proposed_size <= allocation_limit_) {
try {
return upstream_->allocate(bytes, stream);
} catch (...) {
allocated_bytes_ -= proposed_size;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why doesn't this need to be .fetch_sub?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

operator-=(x) is equivalent to fetch_sub(x). fetch_add was used above to retrieve the old value. The old value isn't needed here, thus the less verbose operator-= could be used.

throw;
}
}

return p;
allocated_bytes_ -= proposed_size;
RMM_FAIL("Exceeded memory limit", rmm::bad_alloc);
}

/**
Expand Down Expand Up @@ -165,13 +169,12 @@ class limiting_resource_adaptor final : public device_memory_resource {
* @return true If the two resources are equivalent
* @return false If the two resources are not equal
*/
bool do_is_equal(device_memory_resource const& other) const noexcept override
[[nodiscard]] bool do_is_equal(device_memory_resource const& other) const noexcept override
{
if (this == &other)
return true;
else {
limiting_resource_adaptor<Upstream> const* cast =
dynamic_cast<limiting_resource_adaptor<Upstream> const*>(&other);
auto const* cast = dynamic_cast<limiting_resource_adaptor<Upstream> const*>(&other);
if (cast != nullptr)
return upstream_->is_equal(*cast->get_upstream());
else
Expand All @@ -187,7 +190,8 @@ class limiting_resource_adaptor final : public device_memory_resource {
* @param stream Stream on which to get the mem info.
* @return std::pair contaiing free_size and total_size of memory
*/
std::pair<std::size_t, std::size_t> do_get_mem_info(cuda_stream_view stream) const override
[[nodiscard]] std::pair<std::size_t, std::size_t> do_get_mem_info(
cuda_stream_view stream) const override
{
return {allocation_limit_ - allocated_bytes_, allocation_limit_};
}
Expand Down Expand Up @@ -220,5 +224,4 @@ limiting_resource_adaptor<Upstream> make_limiting_adaptor(Upstream* upstream,
return limiting_resource_adaptor<Upstream>{upstream, allocation_limit};
}

} // namespace mr
} // namespace rmm
} // namespace rmm::mr