Rnn.hpp

Go to the documentation of this file.

// Copyright (c) 2021 Graphcore Ltd. All rights reserved.
#ifndef popnn_Rnn_hpp
#define popnn_Rnn_hpp
 
#include <cassert>
#include <cstdint>
#include <poplar/Graph.hpp>
#include <poplar/Program.hpp>
#include <poplar/Tensor.hpp>
#include <poputil/DebugInfo.hpp>
 
namespace popnn {
namespace rnn {
 
struct RnnParams {
  poplar::Type dataType;
 
  std::size_t batchSize;
 
  std::size_t maxTimeSteps;
 
  std::size_t timeSteps;
 
  poplar::Tensor varTimeSteps;
 
  std::vector<std::size_t> layerSizes;
 
  RnnParams(poplar::Type dataType, std::size_t batchSize, std::size_t timeSteps,
            std::vector<std::size_t> layerSizes);
 
  RnnParams(poplar::Type dataType, std::size_t batchSize,
            std::size_t maxTimeSteps, const poplar::Tensor &varTimeSteps,
            std::vector<std::size_t> layerSizes);
 
  // Return the maximum number of shards.
  std::size_t getMaxShards(const poplar::Graph &graph) const;
 
  // Return the number of bytes of the input per tile.
  std::size_t getInputBytesPerTile(const poplar::Graph &graph) const;
 
  // Return the number of bytes of the output per tile.
  std::size_t getOutputBytesPerTile(const poplar::Graph &graph) const;
 
  // Indicate that time steps are determined by tensor variable.
  bool variableTimeSteps() const;
 
  // Indicate that time steps are determined by tensor variable for each batch.
  bool batchVariableTimeSteps() const;
};
 
poplar::Tensor
createInitialState(poplar::Graph &graph, const RnnParams &params, bool isOutput,
                   unsigned multiple, unsigned numShards,
                   const poplar::DebugContext &debugContext = {});
 
poplar::Tensor
createRecurrentTensor(poplar::Graph &graph, const RnnParams &params,
                      unsigned size, unsigned numShards,
                      const poplar::DebugContext &debugContext = {});
 
poplar::Tensor createInputTensor(poplar::Graph &graph, const RnnParams &params,
                                 unsigned numShards,
                                 const poplar::DebugContext &debugContext = {});
 
poplar::Tensor
createOutputTensor(poplar::Graph &graph, const RnnParams &params,
                   unsigned numShards,
                   const poplar::DebugContext &debugContext = {});
 
poplar::Tensor
createOutputTensor(poplar::Graph &graph, const RnnParams &params,
                   unsigned multiple, unsigned numShards,
                   const poplar::DebugContext &debugContext = {});
 
poplar::Tensor shiftRnnTensor(poplar::Graph &graph, const RnnParams &params,
                              const poplar::Tensor &tBase,
                              const poplar::Tensor &tSingle,
                              poplar::program::Sequence &prog,
                              unsigned numShards,
                              const poplar::DebugContext &debugContext = {});
 
struct RnnSlice {
  std::vector<poplar::Tensor> inputs;
  poplar::Tensor interimIn;
  poplar::Tensor interimOut;
  std::vector<poplar::Tensor> outputs;
};
 
/* Flags set per batch if the current step is within the batchwise step limit.
 * The component tensor(s) are of type `dataType` and shape [`batchSize`].
 */
struct RnnBatchwiseFlags {
  poplar::Tensor mask;
  poplar::Tensor inverse;
 
  bool valid() const { return mask.valid(); };
};
 
struct TimeStepState {
  poplar::Tensor begin;
  poplar::Tensor counter;
  poplar::Tensor variableSeqFlag;
};
 
using LoopBodyType = std::function<poplar::program::Sequence(
    poplar::Graph &graph, const TimeStepState &time, const RnnBatchwiseFlags &,
    std::vector<poplar::Tensor> &, const RnnSlice &slice,
    std::vector<poplar::Tensor> &, poplar::program::Sequence *,
    const poplar::DebugNameAndId &)>;
 
using GatherBodyType = std::function<poplar::program::Sequence(
    poplar::Graph &graph, const RnnSlice &slice, unsigned stepsPerGather,
    poplar::program::Sequence *, const poplar::DebugNameAndId &)>;
 
struct StateSequence {
  poplar::Tensor output;
  std::size_t stateIndex;
};
 
std::vector<poplar::Tensor>
Rnn(poplar::Graph &graph, const RnnParams &params, bool reverse,
    const std::vector<poplar::Tensor> &initState,
    const StateSequence &stateSequence,
    const std::vector<poplar::Tensor> &inputs, const poplar::Tensor *interimIn,
    poplar::Tensor *interimOut, const std::vector<poplar::Tensor> &outputs,
    const std::vector<poplar::Tensor> &created, poplar::program::Sequence &prog,
    const LoopBodyType &loopFn, unsigned numShards,
    poplar::OptionFlags &options,
    const poplar::DebugContext &debugContext = {});
 
std::vector<poplar::Tensor>
Rnn(poplar::Graph &graph, const RnnParams &params,
    const std::vector<poplar::Tensor> &initState,
    const StateSequence &stateSequence,
    const std::vector<poplar::Tensor> &inputs, const poplar::Tensor &interimIn,
    const unsigned numTemps, poplar::program::Sequence &prog,
    const LoopBodyType &loopFn, const std::vector<poplar::Tensor> &gatherInputs,
    const GatherBodyType &gatherFn, unsigned numShards, unsigned stepsPerGather,
    poplar::OptionFlags &options,
    const poplar::DebugContext &debugContext = {});
 
} // namespace rnn
} // namespace popnn
 
#endif // #ifndef popnn_Rnn_hpp