StreamCallback.hpp

// Copyright (c) 2019 Graphcore Ltd. All rights reserved.
 
#ifndef poplar_StreamCallback_hpp
#define poplar_StreamCallback_hpp
 
#include <poplar/CallbackTraits.hpp>
#include <poplar/Quarter.hpp>
#include <poplar/SSOPointer.hpp>
#include <poplar/exceptions.hpp>
 
#include <memory>
 
namespace poplar {
 
// Stream callback interface base class for all data types.
class StreamCallbackBase {
public:
  enum class Result { Success, NotAvailable };
 
  virtual ~StreamCallbackBase() = default;
 
  virtual void complete(){};
 
  virtual void invalidatePrefetched(){};
};
 
class StreamCallback : public StreamCallbackBase {
public:
  virtual ~StreamCallback() = default;
 
  virtual Result prefetch(void *p) = 0;
 
  virtual void fetch(void *) = 0;
};
 
class StreamCallbackWithMetadata : public StreamCallback {
public:
  virtual ~StreamCallbackWithMetadata() = default;
 
  QuarterMetadata getMetadata() const { return metadata; };
 
  void setMetadata(QuarterMetadata md) { metadata = md; };
 
private:
  QuarterMetadata metadata;
};
 
class ResumableStreamCallback : public StreamCallback {
public:
  ResumableStreamCallback();
  ~ResumableStreamCallback();
 
  void notify();
 
  bool isAwaiting() const;
 
protected:
  void wait();
 
private:
  struct ConditionVar;
  SSOPointer<ConditionVar> conditionVar;
};
 
class LegacyStreamCallback : public StreamCallback {
public:
  virtual Result prefetch(void *) final override {
    return Result::NotAvailable;
  }
  virtual void invalidatePrefetched() final override {}
  virtual void complete() final override {}
};
 
class StreamCallbackHandle {
public:
  template <class CallbackImpl,
            typename = typename std::enable_if<
                std::is_base_of<StreamCallback, CallbackImpl>::value ||
                std::is_base_of<StreamCallbackWithMetadata,
                                CallbackImpl>::value>::type>
  StreamCallbackHandle(std::unique_ptr<CallbackImpl> f)
      : callback(std::move(f)) {
    if (!callback) {
      throw poplar_error("Invalid null stream callback");
    }
  }
 
  template <class F, typename = typename std::enable_if<
                         traits::is_callback<F>::value>::type>
  StreamCallbackHandle(F &&f) : callback(makeCallback(std::forward<F>(f))) {}
 
  // Non copy constructible
  StreamCallbackHandle(const StreamCallbackHandle &) = delete;
 
  // Move constructible
  StreamCallbackHandle(StreamCallbackHandle &&) = default;
 
  // This forces the user to provide a non-empty callback and allows the
  // internal implementation to only make use of std::unique_ptr<StreamCallback>
  // if necessary.
  operator std::unique_ptr<StreamCallbackBase>() && {
    return std::move(callback);
  }
 
  operator std::unique_ptr<StreamCallback>() &&;
 
  operator std::unique_ptr<StreamCallbackWithMetadata>() &&;
 
private:
  template <class F>
  static std::unique_ptr<StreamCallback> makeCallback(F &&f) {
    struct NonPrefetchable final : LegacyStreamCallback {
      using Function = typename traits::remove_cvref<F>::type;
 
      NonPrefetchable(F &&f) : function(std::forward<F>(f)) {}
      void fetch(void *p) override { function(p); }
      Function function;
    };
    return std::unique_ptr<NonPrefetchable>(
        new NonPrefetchable(std::forward<F>(f)));
  }
 
  std::unique_ptr<StreamCallbackBase> callback;
};
 
} // end namespace poplar
 
#endif // poplar_StreamCallback_hpp