Graph.hpp

// Copyright (c) 2015 Graphcore Ltd. All rights reserved.
 
#ifndef poplar_Graph_hpp
#define poplar_Graph_hpp
 
#include <poplar/CodeletFileType.hpp>
#include <poplar/DataStream.hpp>
#include <poplar/DebugContext.hpp>
#include <poplar/Device.hpp>
#include <poplar/Executable.hpp>
#include <poplar/FunctionBufferMappingType.hpp>
#include <poplar/GraphElements.hpp>
#include <poplar/Interval.hpp>
#include <poplar/LateInitCallback.hpp>
#include <poplar/PerfEstimateFunc.hpp>
#include <poplar/Program.hpp>
#include <poplar/ReplicatedStreamMode.hpp>
#include <poplar/SerializationFormat.hpp>
#include <poplar/StringRef.hpp>
#include <poplar/Target.hpp>
#include <poplar/Tensor.hpp>
#include <poplar/TensorCloneMethod.hpp>
#include <poplar/TensorRearranger.hpp>
#include <poplar/Type.hpp>
#include <poplar/TypeTraits.hpp>
#include <poplar/VariableMappingMethod.hpp>
#include <poplar/VariableRef.hpp>
#include <poplar/VertexIntrospector.hpp>
#include <poplar/replication_factor.hpp>
 
#include <cstring>
#include <functional>
#include <iosfwd>
#include <memory>
#include <numeric>
#include <string>
#include <vector>
 
namespace poplar {
namespace core {
class GraphBuilder;
}
namespace program {
class Program;
}
 
StringRef versionString();
StringRef packageHash();
 
class Graph {
public:
  Graph();
 
  Graph(const Target &target, replication_factor r = replication_factor(1));
 
  Graph(const Device &device, replication_factor r = replication_factor(1));
  Graph(Graph &&) noexcept;
  Graph &operator=(Graph &&) noexcept;
  ~Graph();
 
  const Target &getTarget() const;
 
  bool addCodelets(StringRef src, CodeletFileType type = CodeletFileType::Auto,
                   StringRef compileFlags = "", StringRef targetName = "");
 
  bool addCodelets(StringRef src, CodeletFileType type, StringRef compileFlags,
                   std::ostream &compileOutput, StringRef targetName = "");
 
  bool addCodelets(ArrayRef<std::string> xs, StringRef compileFlags = "",
                   StringRef targetName = "") {
    bool allAdded = true;
    for (const auto &x : xs)
      allAdded &=
          addCodelets(x, CodeletFileType::Auto, compileFlags, targetName);
    return allAdded;
  }
 
  void addCodelets(std::stringstream &stream, StringRef compileFlags = "",
                   CodeletFileType type = CodeletFileType::CppSource,
                   StringRef targetName = "");
  void addCodelets(std::stringstream &stream, StringRef compileFlags,
                   std::ostream &compileOutput,
                   CodeletFileType type = CodeletFileType::CppSource,
                   StringRef targetName = "");
 
  void addCodeletPlaceholder(std::istream &stream, StringRef targetName = "");
 
  class ConnectionDesc {
    enum class Kind { TensorEdge, ValueEdge, VectorTensorEdge };
    Kind kind;
    std::string field;
    std::vector<Tensor> ts;
    std::unique_ptr<char[]> val;
    TypeTraits traits;
    void connect(Graph &g, const VertexRef &v) const {
      switch (kind) {
      case Kind::TensorEdge:
        g.connect(v[field], ts[0]);
        break;
      case Kind::ValueEdge:
        g.connect(v[field], val.get(), traits);
        break;
      case Kind::VectorTensorEdge:
        g.connect(v[field], ts);
      }
    }
 
  public:
    ConnectionDesc(StringRef field, Tensor t)
        : kind(Kind::TensorEdge), field(field.begin(), field.end()),
          ts({std::move(t)}), traits() {}
 
    ConnectionDesc(StringRef field, ArrayRef<Tensor> tsArr)
        : kind(Kind::VectorTensorEdge), field(field.begin(), field.end()),
          ts(tsArr.begin(), tsArr.end()), traits() {}
 
    template <typename T>
    ConnectionDesc(StringRef field, T v,
                   typename std::enable_if<TypeTraits::isSimpleType<T>()>::type
                       * = nullptr)
        : kind(Kind::ValueEdge), field(field.begin(), field.end()),
          traits(TypeTraits::make<T>()) {
      val.reset(new char[traits.size]);
      std::memcpy(&val[0], &v, traits.size);
    }
 
    friend class Graph;
  };
 
  VertexRef addVertex(ComputeSet cs, StringRef vertexType);
 
  VertexRef addVertex(ComputeSet cs, StringRef vertexType,
                      ArrayRef<ConnectionDesc> connections) {
    auto v = addVertex(cs, vertexType);
    for (const auto &c : connections) {
      c.connect(*this, v);
    }
    return v;
  }
 
  VertexRef addExternalExchangeVertex(ComputeSet cs, StringRef vertexType,
                                      unsigned incomingDownCount,
                                      bool usesEastEdge, bool sendsXReq);
 
  Tensor addVariable(const Type &type, ArrayRef<std::size_t> shape,
                     const DebugContext &debugContext = {});
 
  Tensor addVariable(const Type &type, const Tensor &metadata,
                     ArrayRef<std::size_t> shape,
                     const DebugContext &debugContext = {});
 
  Tensor addVariable(const Type &type, ArrayRef<std::size_t> shape,
                     VariableMappingMethod mappingMethod,
                     const DebugContext &debugContext = {});
 
  Tensor addVariable(const Type &type, const Tensor &metadata,
                     ArrayRef<std::size_t> shape,
                     VariableMappingMethod mappingMethod,
                     const DebugContext &debugContext = {});
 
  Tensor addExternalVariable(const Type &type, ArrayRef<std::size_t> shape,
                             StringRef symbol,
                             const DebugContext &debugContext = {});
 
  template <typename T>
  Tensor addConstant(const Type &type, const Tensor &metadata,
                     ArrayRef<std::size_t> shape, ArrayRef<T> values,
                     const DebugContext &debugContext = {"<const>"}) {
    const auto elements = std::accumulate(
        shape.begin(), shape.end(), std::size_t(1), std::multiplies<size_t>());
    if (values.size() != elements) {
      throw poplar_error("The number of constant values (" +
                         std::to_string(values.size()) +
                         ") does not match "
                         "the number of Tensor elements (" +
                         std::to_string(elements) + ")");
    }
    return addConstant(type, metadata, shape,
                       reinterpret_cast<const void *>(values.data()),
                       TypeTraits::make<T>(), false, debugContext);
  }
 
  template <typename T>
  Tensor addConstant(const Type &type, ArrayRef<std::size_t> shape,
                     ArrayRef<T> values,
                     const DebugContext &debugContext = {"<const>"}) {
    return addConstant(type, {}, shape, values, debugContext);
  }
 
  template <typename T>
  Tensor
  addConstant(const Type &type, const Tensor &metadata,
              ArrayRef<std::size_t> shape, T val,
              const DebugContext &debugContext = {"<const>"},
              typename std::enable_if<TypeTraits::isSimpleType<T>()>::type * =
                  nullptr) {
    return addConstant(type, metadata, shape,
                       reinterpret_cast<const void *>(&val),
                       TypeTraits::make<T>(), true, debugContext);
  }
 
  template <typename T>
  Tensor
  addConstant(const Type &type, ArrayRef<std::size_t> shape, T val,
              const DebugContext &debugContext = {"<const>"},
              typename std::enable_if<TypeTraits::isSimpleType<T>()>::type * =
                  nullptr) {
    return addConstant(type, {}, shape, val, debugContext);
  }
 
  template <typename T>
  Tensor
  addConstant(const Type &type, ArrayRef<std::size_t> shape, const T *val,
              const DebugContext &debugContext = {"<const>"},
              typename std::enable_if<TypeTraits::isSimpleType<T>()>::type * =
                  nullptr) {
    return addConstant(type, shape, static_cast<const void *>(val),
                       TypeTraits::make<T>(), false, debugContext);
  }
 
  Tensor addConstant(const Type &type, ArrayRef<std::size_t> shape,
                     const void *val, const TypeTraits &traits, bool broadcast,
                     const DebugContext &debugContext = {"<const>"});
 
  Tensor addConstant(const Type &type, const Tensor &metadata,
                     ArrayRef<std::size_t> shape, const void *val,
                     const TypeTraits &traits, bool broadcast,
                     const DebugContext &debugContext = {"<const>"});
 
  Tensor addConstantHalf(const Type &type, ArrayRef<std::size_t> shape,
                         uint16_t val,
                         const DebugContext &debugContext = {"<const>"}) {
    return addConstant(type, shape, static_cast<const void *>(&val),
                       TypeTraits::make<IeeeHalf>(), true, debugContext);
  }
 
  Tensor addConstantHalf(const Type &type, ArrayRef<std::size_t> shape,
                         const uint16_t *val,
                         const DebugContext &debugContext = {"<const>"}) {
    return addConstant(type, shape, static_cast<const void *>(val),
                       TypeTraits::make<IeeeHalf>(), false, debugContext);
  }
 
  Tensor clone(const Type &type, const Tensor &t,
               const DebugContext &debugContext = {},
               TensorCloneMethod method =
                   TensorCloneMethod::PRESERVE_ORDER_UNLESS_ALIASES);
 
  Tensor clone(const Type &type, const Tensor &metadata, const Tensor &t,
               const DebugContext &debugContext = {},
               TensorCloneMethod method =
                   TensorCloneMethod::PRESERVE_ORDER_UNLESS_ALIASES);
 
  Tensor cloneN(const Type &type, const Tensor &t, std::size_t N,
                const DebugContext &debugContext = {},
                TensorCloneMethod method =
                    TensorCloneMethod::PRESERVE_ORDER_UNLESS_ALIASES,
                TensorCloneDuplicationMethod duplicationMethod =
                    TensorCloneDuplicationMethod::DUPLICATE_BY_OUTER_DIMENSION);
 
  Tensor cloneN(const Type &type, const Tensor &metadata, const Tensor &t,
                std::size_t N, const DebugContext &debugContext = {},
                TensorCloneMethod method =
                    TensorCloneMethod::PRESERVE_ORDER_UNLESS_ALIASES,
                TensorCloneDuplicationMethod duplicationMethod =
                    TensorCloneDuplicationMethod::DUPLICATE_BY_OUTER_DIMENSION);
 
  Tensor clone(const Tensor &t, const DebugContext &debugContext = {},
               TensorCloneMethod method =
                   TensorCloneMethod::PRESERVE_ORDER_UNLESS_ALIASES) {
    return clone(t.elementType(), t, debugContext, method);
  }
 
  Tensor clone(const Tensor &metadata, const Tensor &t,
               const DebugContext &debugContext = {},
               TensorCloneMethod method =
                   TensorCloneMethod::PRESERVE_ORDER_UNLESS_ALIASES) {
    return clone(t.elementType(), metadata, t, debugContext, method);
  }
 
  Tensor
  cloneN(const Tensor &t, std::size_t N, const DebugContext &debugContext = {},
         TensorCloneMethod method =
             TensorCloneMethod::PRESERVE_ORDER_UNLESS_ALIASES,
         TensorCloneDuplicationMethod duplicationMethod =
             TensorCloneDuplicationMethod::DUPLICATE_BY_OUTER_DIMENSION) {
    return cloneN(t.elementType(), t, N, debugContext, method,
                  duplicationMethod);
  }
 
  Tensor
  cloneN(const Tensor &metadata, const Tensor &t, std::size_t N,
         const DebugContext &debugContext = {},
         TensorCloneMethod method =
             TensorCloneMethod::PRESERVE_ORDER_UNLESS_ALIASES,
         TensorCloneDuplicationMethod duplicationMethod =
             TensorCloneDuplicationMethod::DUPLICATE_BY_OUTER_DIMENSION) {
    return cloneN(t.elementType(), metadata, t, N, debugContext, method,
                  duplicationMethod);
  }
 
  void connect(FieldRef field, const Tensor &tensor);
 
  template <typename T>
  void connect(FieldRef field, T v,
               typename std::enable_if<TypeTraits::isSimpleType<T>()>::type * =
                   nullptr) {
    connect(field, &v, TypeTraits::make<T>());
  }
 
  void connect(FieldRef field, ArrayRef<Tensor> tensors) {
    setFieldSize(field, tensors.size());
    for (unsigned i = 0; i < tensors.size(); ++i) {
      connect(field[i], tensors[i]);
    }
  }
 
  void setPerfEstimate(const VertexRef &v, std::uint64_t cycles,
                       std::uint64_t flops = 0);
 
  void setPerfEstimate(const VertexRef &v, const VertexPerfEstimate &estimate);
 
  VertexPerfEstimate getPerfEstimate(const VertexRef &v) const;
 
  void registerPerfEstimator(StringRef vertexTypeName, PerfEstimateFunc f);
 
  unsigned getNumVertices(void) const;
 
  ComputeSet addComputeSet(const DebugContext &debugContext = {});
 
  void setFieldSize(FieldRef field, std::size_t size);
 
  std::size_t getFieldSize(FieldRef field) const;
 
  std::size_t getMaxFieldDim(StringRef vertexName, StringRef fieldName,
                             unsigned dimIndex) const;
 
  double getMaxVertexFieldValue(StringRef vertexName,
                                StringRef fieldName) const;
 
  template <typename T>
  void setInitialValue(
      FieldRef field, T val,
      typename std::enable_if<TypeTraits::isSimpleType<T>()>::type * =
          nullptr) {
    setInitialValue(field, static_cast<const void *>(&val),
                    TypeTraits::make<T>());
  }
 
  template <typename T>
  void setInitCallback(
      FieldRef field, LateInitCallback<T> callback,
      typename std::enable_if<std::is_arithmetic<T>::value>::type * = nullptr) {
    setInitCallback<T>(field, callback, TypeTraits::make<T>());
  }
 
  // Work arounds for template overload resolution limitations:
  // explicitly convert to ArrayRef ourselves:
  template <typename T>
  void setInitialValue(FieldRef field, const std::vector<T> &v) {
    setInitialValue(field, ArrayRef<T>(v));
  }
  template <typename T>
  void setInitialValue(FieldRef field, const std::initializer_list<T> &l) {
    setInitialValue(field, ArrayRef<T>(l));
  }
 
  void setInitialValueHalf(FieldRef field, uint16_t val) {
    setInitialValue(field, static_cast<const void *>(&val),
                    TypeTraits::make<IeeeHalf>());
  }
 
  template <typename T> void setInitialValue(FieldRef field, ArrayRef<T> val) {
    setFieldSize(field, val.size());
 
    for (unsigned i = 0; i < val.size(); ++i) {
      setInitialValue<T>(field[i], val[i]);
    }
  }
 
  void setInitialValueHalf(FieldRef field, ArrayRef<uint16_t> val) {
    setFieldSize(field, val.size());
 
    for (unsigned i = 0; i < val.size(); ++i) {
      setInitialValueHalf(field[i], val[i]);
    }
  }
 
  template <typename T>
  void setInitialValue(
      const Tensor &t, T val,
      typename std::enable_if<TypeTraits::isSimpleType<T>()>::type * =
          nullptr) {
    setInitialValue(t, static_cast<const void *>(&val), TypeTraits::make<T>());
  }
 
  void setInitialValue(const Tensor &t,
                       const std::map<unsigned, unsigned> &vals);
 
  template <typename T>
  void setInitialValue(const Tensor &t, ArrayRef<T> values) {
    Tensor f = t.flatten();
    if (f.numElements() != values.size()) {
      throw poplar_error("The number of values (" +
                         std::to_string(values.size()) +
                         ") does not match "
                         "the number of Tensor elements (" +
                         std::to_string(f.numElements()) + ")");
    }
    setInitialValue(f, static_cast<const void *>(values.data()),
                    TypeTraits::make<T>());
  }
 
  void setInitialValueHalf(const Tensor &t, uint16_t val) {
    setInitialValue(t, static_cast<const void *>(&val),
                    TypeTraits::make<IeeeHalf>());
  }
 
  void setInitialValueHalf(const Tensor &t, ArrayRef<uint16_t> values) {
    Tensor f = t.flatten();
    if (f.numElements() != values.size()) {
      throw poplar_error("The number of values (" +
                         std::to_string(values.size()) +
                         ") does not match "
                         "the number of Tensor elements (" +
                         std::to_string(f.numElements()) + ")");
    }
    setInitialValue(f, static_cast<const void *>(values.data()),
                    TypeTraits::make<IeeeHalf>());
  }
 
  void createHostWrite(StringRef handle, const Tensor &t,
                       bool rearrangeOnHost = false);
 
  void createHostRead(StringRef handle, const Tensor &t,
                      bool rearrangeOnHost = false);
 
  DataStream addHostToDeviceFIFO(
      StringRef handle, const Type &elementType, std::size_t numElements,
      ReplicatedStreamMode replicatedMode = ReplicatedStreamMode::REPLICATE,
      const OptionFlags &options = {});
 
  DataStream addDeviceToHostFIFO(StringRef handle, const Type &elementType,
                                 std::size_t numElements,
                                 const OptionFlags &options = {});
 
  RemoteBuffer addRemoteBuffer(StringRef handle, const Type &elementType,
                               std::size_t numElements, std::size_t repeats = 1,
                               bool rearrangeOnHost = false,
                               bool optimiseMemory = false);
 
  void outputVertexGraph(std::ostream &outputStream,
                         ArrayRef<program::Program> progs = {}) const;
 
  void outputComputeGraph(std::ostream &outputStream,
                          ArrayRef<program::Program> progs = {}) const;
 
  void setTileMapping(VertexRef v, unsigned tileNum);
 
  void setTileMapping(const Tensor &t, unsigned tileNum);
 
  using TileToTensorMapping = std::vector<std::vector<Interval>>;
 
  TileToTensorMapping getTileMapping(const Tensor &t,
                                     bool requireComplete = true,
                                     bool allowExternal = false) const;
 
  TileToTensorMapping getTileMapping(const Tensor &t, bool *isComplete,
                                     bool allowExternal = false) const;
 
  TileToTensorMapping getVariableTileMapping(const Tensor &t) const;
 
  void setTileMapping(const Tensor &t, const TileToTensorMapping &mapping);
 
  Tensor getVariable(VariableRef v) const;
 
  bool isConstant(VariableRef v) const;
 
  std::vector<std::vector<Interval>>
  getSortedContiguousRegions(const Tensor &t, ArrayRef<Interval> regions,
                             bool removeAliasedIntervals = false,
                             std::vector<std::size_t> *aliases = nullptr) const;
 
  void reorderToSimplify(Tensor *t, ArrayRef<Tensor *> ts,
                         bool requireSimplestOrder = true) const;
 
  TensorRearranger getSimplifyingRearranger(const Tensor &t) const;
 
  Tensor findUnbroadcastTensor(const Tensor &t) const;
 
  void serializeTensors(std::ostream &out, ArrayRef<Tensor> tensors,
                        SerializationFormat format) const;
 
  std::vector<Tensor> deserializeTensors(std::istream &in,
                                         SerializationFormat format);
 
  Graph createVirtualGraph(unsigned numTilesPerIPU);
 
  Graph createVirtualGraph(unsigned lowerTile, unsigned upperTile);
 
  Graph createVirtualGraph(const std::vector<unsigned> &perIpuTiles);
 
  Graph getTopLevelGraph();
 
  unsigned getReplicationFactor() const;
 
  Tensor addReplicationIndexConstant(const DebugContext &debugContext = {});
 
  void serialize(std::ostream &out, SerializationFormat format) const;
 
  void serialize(std::ostream &out, ArrayRef<program::Program> progs,
                 SerializationFormat format) const;
 
  Function addFunction(const program::Program &program);
 
  FunctionBuffer addFunctionBuffer(const Function &f,
                                   FunctionBufferMappingType mappingType);
 
  using HostFunctionArgument = std::tuple<Type, std::size_t>;
  HostFunction addHostFunction(StringRef handle,
                               ArrayRef<HostFunctionArgument> inputs,
                               ArrayRef<HostFunctionArgument> outputs);
 
  unsigned convertTileToTopLevelGraphTile(unsigned tileId) const;
 
  unsigned convertTopLevelGraphTileToTile(unsigned topLevelGraphTileId) const;
 
  unsigned convertVirtualTileToPhysicalTile(unsigned tileId) const;
 
  unsigned convertPhysicalTileToVirtualTile(unsigned physicalTileId) const;
 
  unsigned convertPhysicalTileToVirtualTile(unsigned ipuId,
                                            unsigned physicalTileId) const;
 
  bool hasCodelet(StringRef codeletName) const;
 
  // Implementation
  Graph(std::unique_ptr<core::GraphBuilder>, Target target);
  const core::GraphBuilder &getImpl() const { return *impl; }
  core::GraphBuilder &getImpl() { return *impl; }
 
private:
  std::unique_ptr<core::GraphBuilder> impl;
  Target target;
 
  void setInitialValue(FieldRef field, const void *val, const TypeTraits &);
  template <typename T>
  void setInitCallback(FieldRef field, LateInitCallback<T> callback,
                       const TypeTraits &);
  void setInitialValue(const Tensor &t, const void *val, const TypeTraits &);
  void connect(FieldRef field, void *val, const TypeTraits &);
  void checkFieldSubgraph(const FieldRef &f) const;
  void checkVertexSubgraph(const VertexRef &v) const;
};
 
} // End namespace poplar.
 
#endif // poplar_Graph_hpp