pytorch/c10/util/Synchronized.h
Dhruv Matani e1eb876ade [PyTorch] Update Synchronized<T>::withLock() to return the type/value from the aceepted callable (#74108)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/74108

Currently, `void c10::Synchronized<T>::withLock(CB)` accepts a callback but returns `void`. For a callback (CB) returning a value, `withLock()` can not be used to retrieve the returned value. This change makes it possible for the caller to retrieve that value. The only gotcha is that if the value has only move semantics (such as `std::unique_ptr<T>`, then this won't work as currently written). To be able to do that, one needs to use `std::enable_if<T>` to enable a separate overload that works only for moveable types (detectable using `std::is_move_constructible<T>`).
ghstack-source-id: 151471745

Test Plan:
Built model tracer using the command:

```
buck build -c pt.disable_per_op_profiling=0 -c pt.enable_record_kernel_dtype=1 --show-output xplat/caffe2/fb/model_tracer:model_tracer
```

Reviewed By: malfet

Differential Revision: D34822915

fbshipit-source-id: 2ade3d963d1e997a791508ec54a930827f7849b5
(cherry picked from commit d5429073137b33978882299c5ee5cc7ffb964f8e)
2022-03-16 16:18:36 +00:00

62 lines
1.9 KiB
C++

#pragma once
#include <mutex>
namespace c10 {
/**
* A very simple Synchronization class for error-free use of data
* in a multi-threaded context. See folly/docs/Synchronized.md for
* the inspiration of this class.
*
* Full URL:
* https://github.com/facebook/folly/blob/main/folly/docs/Synchronized.md
*
* This class implements a small subset of the generic functionality
* implemented by folly:Synchronized<T>. Specifically, only withLock<T>
* is implemeted here since it's the smallest possible API that is
* able to cover a large surface area of functionality offered by
* folly::Synchronized<T>.
*/
template <typename T>
class Synchronized final {
mutable std::mutex mutex_;
T data_;
public:
Synchronized() = default;
Synchronized(T const& data) : data_(data) {}
Synchronized(T&& data) : data_(data) {}
// Don't permit copy construction, move, assignment, or
// move assignment, since the underlying std::mutex
// isn't necessarily copyable/moveable.
Synchronized(Synchronized const&) = delete;
Synchronized(Synchronized&&) = delete;
Synchronized operator=(Synchronized const&) = delete;
Synchronized operator=(Synchronized&&) = delete;
/**
* To use, call withLock<T> with a callback that accepts T either
* by copy or by reference. Use the protected variable in the
* provided callback safely.
*/
template <typename CB>
typename std::result_of<CB(T&)>::type withLock(CB cb) {
std::lock_guard<std::mutex> guard(this->mutex_);
return cb(this->data_);
}
/**
* To use, call withLock<T> with a callback that accepts T either
* by copy or by const reference. Use the protected variable in
* the provided callback safely.
*/
template <typename CB>
typename std::result_of<CB(T const&)>::type withLock(CB cb) const {
std::lock_guard<std::mutex> guard(this->mutex_);
return cb(this->data_);
}
};
} // end namespace c10