diff --git a/paddle/fluid/framework/channel.h b/paddle/fluid/framework/channel.h index 9f8fb12098d622058a86f83c1c42a1feb1cfb2e2..51e2b03f9cb9abf6b3effe4035d4eec2ba4f9fbf 100644 --- a/paddle/fluid/framework/channel.h +++ b/paddle/fluid/framework/channel.h @@ -15,23 +15,43 @@ limitations under the License. */ #pragma once #include // for size_t +#include #include #include "paddle/fluid/platform/enforce.h" namespace paddle { namespace framework { +enum class ChannelAction { + SEND = 0, + RECEIVE = 1, + CLOSE = 2, +}; + // Channel is the abstract class of buffered and un-buffered channels. template class Channel { public: + virtual bool CanSend() = 0; + virtual bool CanReceive() = 0; virtual bool Send(T*) = 0; virtual bool Receive(T*) = 0; virtual size_t Cap() = 0; virtual void Lock() = 0; + virtual void Unlock() = 0; + virtual bool IsClosed() = 0; virtual void Close() = 0; virtual ~Channel() {} + + virtual void AddToSendQ(const void* referrer, T* data, + std::shared_ptr cond, + std::function cb) = 0; + virtual void AddToReceiveQ(const void* referrer, T* data, + std::shared_ptr cond, + std::function cb) = 0; + virtual void RemoveFromSendQ(const void* referrer) = 0; + virtual void RemoveFromReceiveQ(const void* referrer) = 0; }; // Forward declaration of channel implementations. @@ -80,6 +100,27 @@ class ChannelHolder { return channel != nullptr ? channel->Receive(data) : false; } + bool IsClosed() { + if (IsInitialized()) { + return holder_->IsClosed(); + } + return false; + } + + bool CanSend() { + if (IsInitialized()) { + return holder_->CanSend(); + } + return false; + } + + bool CanReceive() { + if (IsInitialized()) { + return holder_->CanReceive(); + } + return false; + } + void close() { if (IsInitialized()) holder_->Close(); } @@ -97,6 +138,50 @@ class ChannelHolder { if (IsInitialized()) holder_->Unlock(); } + template + void AddToSendQ(const void* referrer, T* data, + std::shared_ptr cond, + std::function cb) { + if (IsInitialized()) { + Channel* channel = static_cast*>(holder_->Ptr()); + if (channel != nullptr) { + channel->AddToSendQ(referrer, data, cond, cb); + } + } + } + + template + void AddToReceiveQ(const void* referrer, T* data, + std::shared_ptr cond, + std::function cb) { + if (IsInitialized()) { + Channel* channel = static_cast*>(holder_->Ptr()); + if (channel != nullptr) { + channel->AddToReceiveQ(referrer, data, cond, cb); + } + } + } + + template + void RemoveFromSendQ(const void* referrer) { + if (IsInitialized()) { + Channel* channel = static_cast*>(holder_->Ptr()); + if (channel != nullptr) { + channel->RemoveFromSendQ(referrer); + } + } + } + + template + void RemoveFromReceiveQ(const void* referrer) { + if (IsInitialized()) { + Channel* channel = static_cast*>(holder_->Ptr()); + if (channel != nullptr) { + channel->RemoveFromReceiveQ(referrer); + } + } + } + inline bool IsInitialized() const { return holder_ != nullptr; } inline const std::type_index Type() { @@ -113,6 +198,9 @@ class ChannelHolder { virtual ~Placeholder() {} virtual const std::type_index Type() const = 0; virtual void* Ptr() const = 0; + virtual bool IsClosed() = 0; + virtual bool CanSend() = 0; + virtual bool CanReceive() = 0; virtual void Close() = 0; virtual void Lock() = 0; virtual void Unlock() = 0; @@ -129,6 +217,27 @@ class ChannelHolder { virtual void* Ptr() const { return static_cast(channel_.get()); } + virtual bool IsClosed() { + if (channel_) { + return channel_->IsClosed(); + } + return false; + } + + virtual bool CanSend() { + if (channel_) { + return channel_->CanSend(); + } + return false; + } + + virtual bool CanReceive() { + if (channel_) { + return channel_->CanReceive(); + } + return false; + } + virtual void Close() { if (channel_) channel_->Close(); } diff --git a/paddle/fluid/framework/channel_impl.h b/paddle/fluid/framework/channel_impl.h index a4561031fd8c49613269e7008ce558f25f9765e4..c194c03e264cccfa9ad755ea24bc6372e82bfb00 100644 --- a/paddle/fluid/framework/channel_impl.h +++ b/paddle/fluid/framework/channel_impl.h @@ -29,32 +29,50 @@ class ChannelImpl : public paddle::framework::Channel { friend void paddle::framework::CloseChannel(Channel *); public: + virtual bool CanSend(); + virtual bool CanReceive(); virtual bool Send(T *); virtual bool Receive(T *); virtual size_t Cap() { return cap_; } virtual void Lock(); virtual void Unlock(); + virtual bool IsClosed(); virtual void Close(); - ChannelImpl(size_t); virtual ~ChannelImpl(); + virtual void AddToSendQ(const void *referrer, T *data, + std::shared_ptr cond, + std::function cb); + virtual void AddToReceiveQ(const void *referrer, T *data, + std::shared_ptr cond, + std::function cb); + + virtual void RemoveFromSendQ(const void *referrer); + virtual void RemoveFromReceiveQ(const void *referrer); + private: struct QueueMessage { T *data; - std::condition_variable_any cond; + std::shared_ptr cond; bool chan_closed = false; bool completed = false; + const void *referrer; // TODO(thuan): figure out better way to do this + std::function callback; - QueueMessage(T *item) : data(item) {} + QueueMessage(T *item) + : data(item), cond(std::make_shared()) {} + + QueueMessage(T *item, std::shared_ptr cond) + : data(item), cond(cond) {} void Wait(std::unique_lock &lock) { - cond.wait(lock, [this]() { return completed; }); + cond->wait(lock, [this]() { return completed; }); } void Notify() { completed = true; - cond.notify_all(); + cond->notify_all(); } }; @@ -87,6 +105,18 @@ ChannelImpl::ChannelImpl(size_t capacity) PADDLE_ENFORCE_GE(capacity, 0); } +template +bool ChannelImpl::CanSend() { + std::lock_guard lock{mu_}; + return !closed_ && (!recvq.empty() || buf_.size() < cap_); +} + +template +bool ChannelImpl::CanReceive() { + std::lock_guard lock{mu_}; + return !(closed_ && buf_.empty()) && (!sendq.empty() || buf_.size() > 0); +} + template bool ChannelImpl::Send(T *item) { send_ctr++; @@ -105,7 +135,24 @@ bool ChannelImpl::Send(T *item) { std::shared_ptr m = recvq.front(); recvq.pop_front(); // Do the data transfer - *(m->data) = std::move(*item); + // We will do this data transfer if either of the following + // cases are true + // 1. callback == nullptr // This means it was a regular channel send + // 2. callback returns true + bool do_send = true; + if (m->callback != nullptr) do_send = m->callback(ChannelAction::SEND); + if (do_send) + *(m->data) = std::move(*item); + else + // We cannot do the data transfer because + // this QueueMessage was added by Select + // and some other case was executed. + // So call the Send function again. + // We do not care about notifying other + // because they would have been notified + // by the executed select case. + return Send(item); + // Wake up the blocked process and unlock m->Notify(); lock.unlock(); @@ -150,7 +197,25 @@ bool ChannelImpl::Receive(T *item) { std::shared_ptr m = sendq.front(); sendq.pop_front(); // Do the data transfer - *item = std::move(*(m->data)); + // We will do this data transfer if either of the following + // cases are true + // 1. callback == nullptr // This means it was a regular channel send + // 2. callback returns true + bool do_receive = true; + if (m->callback != nullptr) + do_receive = m->callback(ChannelAction::RECEIVE); + if (do_receive) + *item = std::move(*(m->data)); + else + // We cannot do the data transfer because + // this QueueMessage was added by Select + // and some other case was executed. + // So call the Receive function again. + // We do not care about notifying other + // because they would have been notified + // by the executed select case. + return Receive(item); + // Wake up the blocked process and unlock m->Notify(); lock.unlock(); @@ -186,6 +251,12 @@ void ChannelImpl::Unlock() { mu_.unlock(); } +template +bool ChannelImpl::IsClosed() { + std::lock_guard lock{mu_}; + return closed_; +} + template void ChannelImpl::Close() { std::unique_lock lock{mu_}; @@ -203,6 +274,12 @@ void ChannelImpl::Close() { std::shared_ptr m = recvq.front(); recvq.pop_front(); m->chan_closed = true; + + // Execute callback function (if any) + if (m->callback != nullptr) { + m->callback(ChannelAction::CLOSE); + } + m->Notify(); } @@ -211,10 +288,72 @@ void ChannelImpl::Close() { std::shared_ptr m = sendq.front(); sendq.pop_front(); m->chan_closed = true; + + // Execute callback function (if any) + if (m->callback != nullptr) { + m->callback(ChannelAction::CLOSE); + } + m->Notify(); } } +template +void ChannelImpl::AddToSendQ( + const void *referrer, T *data, + std::shared_ptr cond, + std::function cb) { + std::lock_guard lock{mu_}; + auto m = std::make_shared(data, cond); + m->referrer = referrer; + m->callback = cb; + sendq.push_back(m); +} + +template +void ChannelImpl::AddToReceiveQ( + const void *referrer, T *data, + std::shared_ptr cond, + std::function cb) { + std::lock_guard lock{mu_}; + auto m = std::make_shared(data, cond); + m->referrer = referrer; + m->callback = cb; + recvq.push_back(m); +} + +template +void ChannelImpl::RemoveFromSendQ(const void *referrer) { + std::lock_guard lock{mu_}; + + for (auto it = sendq.begin(); it != sendq.end();) { + std::shared_ptr sendMsg = (std::shared_ptr)*it; + + if (sendMsg->referrer == referrer) { + it = sendq.erase(it); + send_ctr--; + } else { + ++it; + } + } +} + +template +void ChannelImpl::RemoveFromReceiveQ(const void *referrer) { + std::lock_guard lock{mu_}; + + for (auto it = recvq.begin(); it != recvq.end();) { + std::shared_ptr recvMsg = (std::shared_ptr)*it; + + if (recvMsg->referrer == referrer) { + it = recvq.erase(it); + recv_ctr--; + } else { + ++it; + } + } +} + template ChannelImpl::~ChannelImpl() { Close();