HostFunctionCallback.hpp

// Copyright (c) 2021 Graphcore Ltd. All rights reserved.
#ifndef poplar_HostCallback_hpp
#define poplar_HostCallback_hpp
 
#include <poplar/CallbackTraits.hpp>
 
#include <memory>
 
namespace poplar {
 
class HostCallback {
public:
  virtual ~HostCallback();
 
  virtual void operator()(ArrayRef<const void *> inputs,
                          ArrayRef<void *> outputs) = 0;
};
 
class HostCallbackHandle {
public:
  template <class CallbackImpl,
            typename = typename std::enable_if<
                std::is_base_of<HostCallback, CallbackImpl>::value>::type>
  HostCallbackHandle(std::unique_ptr<CallbackImpl> f)
      : callback(std::move(f)) {}
 
  template <class F, typename = typename std::enable_if<
                         traits::is_host_callback<F>::value>::type>
  HostCallbackHandle(F &&f) : callback(makeCallback(std::forward<F>(f))) {}
 
  // Non copy constructible
  HostCallbackHandle(const HostCallbackHandle &) = delete;
 
  // Move constructible
  HostCallbackHandle(HostCallbackHandle &&) = default;
 
  // This forces the user to provide a non-empty callback and allows the
  // internal implementation to only make use of std::unique_ptr<HostCallback>
  // if necessary.
  operator std::unique_ptr<HostCallback>() && { return std::move(callback); }
 
private:
  template <class F> std::unique_ptr<HostCallback> makeCallback(F &&f) {
    struct Callback final : HostCallback {
      using Function = typename traits::remove_cvref<F>::type;
 
      Function function;
 
      Callback(F &&f) : function(std::forward<F>(f)) {}
 
      void operator()(ArrayRef<const void *> ins,
                      ArrayRef<void *> outs) override {
        function(ins, outs);
      }
    };
    return std::unique_ptr<HostCallback>(new Callback{std::forward<F>(f)});
  }
 
  std::unique_ptr<HostCallback> callback;
};
 
} // namespace poplar
 
#endif // poplar_HostCallback_hpp