-
Notifications
You must be signed in to change notification settings - Fork 197
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.
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
fix race condition in limiting resource adapter #869
Conversation
@@ -130,11 +132,12 @@ class limiting_resource_adaptor final : public device_memory_resource { | |||
void* p = nullptr; | |||
|
|||
std::size_t proposed_size = rmm::detail::align_up(bytes, allocation_alignment_); | |||
if (proposed_size + allocated_bytes_ <= allocation_limit_) { | |||
allocated_bytes_ += proposed_size; | |||
if (allocated_bytes_ <= allocation_limit_) { | |||
p = upstream_->allocate(bytes, stream); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
what if allocate throws? should we try/catch?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Hm, good point. We should probably catch the exception and decrement allocated_bytes_
and then rethrow.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done.
allocated_bytes_ += proposed_size; | ||
if (allocated_bytes_ <= allocation_limit_) { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Apart from the fact that this is doing a redundant atomic load, I don't believe this is the semantics of what we want. This allows an intervening thread to impact the value of allocated_bytes_
between the fetch_add
and the load
.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Fixed.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
One question.
try { | ||
return upstream_->allocate(bytes, stream); | ||
} catch (...) { | ||
allocated_bytes_ -= proposed_size; |
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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.
@gpucibot merge |
Fixes #868
Also fixed some clang-tidy warnings.