GraphFunction.hpp

Go to the documentation of this file.

// Copyright (c) 2017 Graphcore Ltd. All rights reserved.
#ifndef poputil_GraphFunction_hpp
#define poputil_GraphFunction_hpp
#include <poplar/DebugContext.hpp>
#include <poplar/Graph.hpp>
#include <poplar/Program.hpp>
 
namespace poputil {
 
namespace graphfn {
 
enum ArgType { InputArg, OutputArg, InOutArg, CreatedArg };
 
struct ArgSig {
  ArgType type;
  poplar::Tensor similarTensor;
  std::string debugName;
  ArgSig(ArgType type, poplar::Tensor tensor, std::string debugName)
      : type(type), similarTensor(std::move(tensor)),
        debugName(std::move(debugName)) {}
};
 
inline ArgSig input(poplar::Tensor similar, std::string debugName = "") {
  return ArgSig(InputArg, std::move(similar), std::move(debugName));
}
 
inline ArgSig inout(poplar::Tensor similar, std::string debugName = "") {
  return ArgSig(InOutArg, std::move(similar), std::move(debugName));
}
 
inline ArgSig output(poplar::Tensor similar, std::string debugName = "") {
  return ArgSig(OutputArg, std::move(similar), std::move(debugName));
}
 
inline ArgSig created(std::string debugName = "") {
  return ArgSig(CreatedArg, poplar::Tensor(), std::move(debugName));
}
 
using Signature = std::vector<ArgSig>;
 
class VoidFunction {
  poplar::Graph &graph;
  Signature sig;
  bool inlined;
  poplar::program::Sequence prog;
  poplar::Function func;
  std::vector<poplar::Tensor> params;
 
public:
  VoidFunction(VoidFunction &&fn);
  VoidFunction(poplar::Graph &graph, Signature sig,
               std::function<void(std::vector<poplar::Tensor> &,
                                  poplar::program::Sequence &)>
                   f,
               bool inlined = false,
               const poplar::DebugContext &debugContext = {});
 
  VoidFunction(poplar::Graph &graph, Signature sig,
               std::function<void(std::vector<poplar::Tensor> &,
                                  poplar::program::Sequence &,
                                  const poplar::DebugNameAndId &)>
                   f,
               bool inlined = false,
               const poplar::DebugContext &debugContext = {});
 
  void operator()(std::vector<poplar::Tensor> &args,
                  poplar::program::Sequence &seq,
                  const poplar::DebugContext &dc = {});
};
 
class ProgramFunction {
  VoidFunction voidFunc;
 
public:
  ProgramFunction(
      poplar::Graph &graph, Signature sig,
      std::function<poplar::program::Program(std::vector<poplar::Tensor> &)> f,
      bool inlined = false, const poplar::DebugContext &debugContext = {});
 
  ProgramFunction(
      poplar::Graph &graph, Signature sig,
      std::function<poplar::program::Program(std::vector<poplar::Tensor> &,
                                             const poplar::DebugNameAndId &)>
          f,
      bool inlined = false, const poplar::DebugContext &debugContext = {});
 
  poplar::program::Program
  operator()(std::vector<poplar::Tensor> &args,
             const poplar::DebugContext &debugContext = {});
};
 
class TensorFunction {
  VoidFunction voidFunc;
 
public:
  TensorFunction(poplar::Graph &graph, Signature sig,
                 std::function<poplar::Tensor(std::vector<poplar::Tensor> &,
                                              poplar::program::Sequence &)>
                     f,
                 bool inlined = false,
                 const poplar::DebugContext &debugContext = {});
 
  TensorFunction(poplar::Graph &graph, Signature sig,
                 std::function<poplar::Tensor(std::vector<poplar::Tensor> &,
                                              poplar::program::Sequence &,
                                              const poplar::DebugNameAndId &)>
                     f,
                 bool inlined = false,
                 const poplar::DebugContext &debugContext = {});
 
  poplar::Tensor operator()(std::vector<poplar::Tensor> &args,
                            poplar::program::Sequence &prog,
                            const poplar::DebugContext &debugContext = {});
};
 
} // namespace graphfn
} // namespace poputil
 
#endif // poputil_GraphFunction_hpp