Program.hpp

// Copyright (c) 2016 Graphcore Ltd. All rights reserved.
// Definitions of program classes.
 
#ifndef poplar_Program_hpp
#define poplar_Program_hpp
 
#include <cstdint>
#include <memory>
#include <ostream>
#include <poplar/DataStream.hpp>
#include <poplar/DebugContext.hpp>
#include <poplar/GraphElements.hpp>
#include <poplar/OptionFlags.hpp>
#include <poplar/StringRef.hpp>
#include <poplar/SyncType.hpp>
#include <poplar/Tensor.hpp>
#include <poplar/TypeTraits.hpp>
 
namespace poplar {
namespace core {
class ProgramImpl;
}
 
namespace program {
 
class Program {
public:
  Program();
  Program(const Program &p);
  Program(Program &&p) noexcept;
  Program &operator=(const Program &p);
  Program &operator=(Program &&p) noexcept;
  virtual ~Program();
 
  // Check if program is empty
  bool isEmpty() const { return impl.get() == nullptr; }
 
  // Implementation
  core::ProgramImpl &getImpl() { return *impl; }
  const core::ProgramImpl &getImpl() const { return *impl; }
 
  friend bool operator==(const Program &lhs, const Program &rhs);
 
protected:
  std::unique_ptr<core::ProgramImpl> impl;
};
 
bool operator==(const Program &lhs, const Program &rhs);
 
class Execute : public Program {
public:
  explicit Execute(ComputeSet cs, const DebugContext &debugContext = {});
 
  Execute(ComputeSet cs, Tensor t, const DebugContext &debugContext = {});
};
 
class Sequence : public Program {
  template <class... T> void add_many(const Program &first, T &&...rest) {
    add(first);
    add_many(rest...);
  }
  void add_many() {}
  void init();
  void init(const DebugContext &debugContext);
 
public:
  Sequence(const DebugContext &debugContext = {}) { init(debugContext); }
 
  Sequence(std::initializer_list<Program> programs,
           const DebugContext &debugContext = {}) {
    init(debugContext);
    for (const auto &p : programs) {
      add(p);
    }
  }
 
  void add(const Program &p);
};
 
class Repeat : public Program {
public:
  Repeat(unsigned count, const Program &prog,
         const DebugContext &debugContext = {});
};
 
class RepeatWhileFalse : public Program {
public:
  RepeatWhileFalse(const Program &cond, Tensor predicate, const Program &body,
                   const DebugContext &debugContext = {});
};
 
class RepeatWhileTrue : public Program {
public:
  RepeatWhileTrue(const Program &cond, Tensor predicate, const Program &body,
                  const DebugContext &debugContext = {});
};
 
class Loop : public Program {
public:
  Loop(const Program &prog, const DebugContext &debugContext = {});
};
 
class If : public Program {
public:
  If(Tensor predicate, const Program &trueBody, const Program &falseBody,
     const DebugContext &debugContext = {});
};
 
class Switch : public Program {
  Switch(Tensor control, const Program &defaultCaseBody,
         const bool unreachableDefault, const DebugContext &debugContext = {});
 
public:
  Switch(Tensor control,
         const std::vector<std::pair<std::int32_t, Program>> &cases,
         const DebugContext &debugContext = {});
  Switch(Tensor control,
         const std::vector<std::pair<std::int32_t, Program>> &cases,
         const Program &defaultCaseBody, const DebugContext &debugContext = {});
  Switch(Tensor control, const DebugContext &debugContext = {});
  Switch(Tensor control, const Program &defaultCaseBody,
         const DebugContext &debugContext = {});
 
  Switch &add(std::int32_t value, const Program &body);
 
  static Switch switchWithBoundsChecking(
      Tensor control,
      const std::vector<std::pair<std::int32_t, Program>> &cases,
      const DebugContext &debugContext = {});
  static Switch
  switchWithUnreachableDefault(Tensor control,
                               const DebugContext &debugContext = {});
};
 
class Copy : public Program {
  Copy(const DataStream &stream, Tensor dst, bool rearrangeOnHost,
       Tensor offset, size_t repeats, bool optimiseMemory,
       const OptionFlags &options = {}, const DebugContext &debugContext = {});
  Copy(Tensor src, const DataStream &stream, bool rearrangeOnHost,
       Tensor offset, size_t repeats, bool optimiseMemory,
       const OptionFlags &options = {}, const DebugContext &debugContext = {});
 
public:
  Copy(Tensor src, Tensor dst, bool dontOutline = false,
       const DebugContext &debugContext = {});
 
  // allow long lines for link
  // clang-format off
  // clang-format on
  Copy(const DataStream &stream, Tensor dst, bool optimiseMemory = false,
       const DebugContext &debugContext = {});
 
  // allow long lines for link
  // clang-format off
  // clang-format on
  Copy(Tensor src, const DataStream &stream, bool optimiseMemory = false,
       const DebugContext &debugContext = {});
 
  Copy(const RemoteBuffer &buffer, Tensor dst,
       const DebugContext &debugContext = {});
 
  Copy(const RemoteBuffer &buffer, Tensor dst, Tensor offset,
       const DebugContext &debugContext = {});
 
  Copy(Tensor src, const RemoteBuffer &buffer,
       const DebugContext &debugContext = {});
 
  Copy(Tensor src, const RemoteBuffer &buffer, Tensor offset,
       const DebugContext &debugContext = {});
 
  Copy(const FunctionBuffer &buffer, const Function &function,
       const DebugContext &debugContext = {});
};
 
class CrossReplicaCopy : public Program {
public:
  CrossReplicaCopy(Tensor src, Tensor dst,
                   std::map<unsigned, unsigned> replicaMap,
                   const DebugContext &debugContext = {});
};
 
class WriteUndef : public Program {
public:
  WriteUndef(Tensor t, const DebugContext &debugContext = {});
};
 
class AssumeEqualAcrossReplicas : public Program {
public:
  AssumeEqualAcrossReplicas(Tensor t, const DebugContext &debugContext = {});
};
 
class Block : public Program {
public:
  Block(const Program &p, const DebugContext &debugContext = {});
};
 
class Sync : public Program {
public:
  Sync(SyncType type, const DebugContext &debugContext = {});
};
 
class Call : public Program {
public:
  Call(Function f, const DebugContext &debugContext = {});
 
  //                 arguments
  Call(HostFunction f, ArrayRef<Tensor> inputs, ArrayRef<Tensor> outputs,
       const DebugContext &debugContext = {});
};
 
class PrintTensor : public Program {
public:
  PrintTensor(Tensor t, const DebugContext &debugContext = {});
  PrintTensor(StringRef title, Tensor t, const DebugContext &debugContext = {});
};
 
class ErrorProgram : public Program {
public:
  ErrorProgram(StringRef message, Tensor debugTensor,
               const DebugContext &debugContext = {});
};
 
class Abort : public Program {
public:
  Abort(const DebugContext &debugContext = {});
 
  Abort(const std::string &message, const DebugContext &debugContext = {});
};
 
class AbortOnCondition : public Program {
public:
  AbortOnCondition(Tensor predicate, const DebugContext &debugContext = {});
 
  AbortOnCondition(Tensor predicate, const std::string &message,
                   const DebugContext &debugContext = {});
};
 
void dumpProgram(const Graph &graph, const Program &program, std::ostream &out);
 
} // namespace program
} // namespace poplar
#endif // poplar_Program_hpp