// nnet3/nnet-compile.h
  
  // Copyright 2015-2016    Johns Hopkins University (author: Daniel Povey)
  
  // See ../../COPYING for clarification regarding multiple authors
  //
  // Licensed under the Apache License, Version 2.0 (the "License");
  // you may not use this file except in compliance with the License.
  // You may obtain a copy of the License at
  //
  //  http://www.apache.org/licenses/LICENSE-2.0
  //
  // THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  // KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED
  // WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
  // MERCHANTABLITY OR NON-INFRINGEMENT.
  // See the Apache 2 License for the specific language governing permissions and
  // limitations under the License.
  
  #ifndef KALDI_NNET3_NNET_COMPILE_H_
  #define KALDI_NNET3_NNET_COMPILE_H_
  
  #include "nnet3/nnet-component-itf.h"
  #include "nnet3/nnet-nnet.h"
  #include "nnet3/nnet-parse.h"
  #include "nnet3/nnet-computation.h"
  #include "nnet3/nnet-computation-graph.h"
  
  #include <iostream>
  
  namespace kaldi {
  namespace nnet3 {
  
  struct CompilerOptions {
    bool output_debug_info;
  
    CompilerOptions(): output_debug_info(true) { }
  };
  
  /// This class creates an initial version of the NnetComputation, without any
  /// optimization or sharing of matrices.    Note: for a user-level interface
  /// that includes optimization, see class CachingOptimizingCompiler in
  /// nnet-optimize.h.
  class Compiler {
   public:
    // Constructor that takes one computation request (this is the normal case).
    Compiler(const ComputationRequest &request,
             const Nnet &nnet);
  
    // Constructor with a sequence of computation requests, for multiple
    // computation segments (used when creating online computations).
    Compiler(const std::vector<const ComputationRequest*> &request,
             const Nnet &nnet);
  
    void CreateComputation(const CompilerOptions &opts,
                           NnetComputation *computation);
  
   private:
    // requests_ is the sequence of computation requests, one for each segment; it
    // will contain just one element in the normal case, but more when we're
    // compiling a multi-segment / 'online' computation.
    std::vector<const ComputationRequest*> requests_;
    const Nnet &nnet_;
    ComputationGraph graph_;
  
    // Some generic information about each step of the computation... a step is an
    // instance of a NetworkNode, but a NetworkNode may in general have multiple
    // steps.  A single step may turn into no commands (for input nodes), or
    // multiple commands.  The StepInfo also contains info about the backprop
    // corresponding to its forward command.
    struct StepInfo {
      int32 node_index;  // network-node index
      int32 value;  // sub-matrix index of value that this step outputs.
      int32 deriv;  // sub-matrix index of derivative at the output of this step; zero
                    // if not used (note: index zero is reserved for the empty
                    // matrix).
  
      int32 segment;  // normally 0 except for online/multi-segment computations,
                      // identifies the segment of which this step is a part (each
                      // segment in the sequence has a different
                      // ComputationRequest).
  
      // precomputed_indexes_index is the index into the
      // component_precomputed_indexes array in the NnetComputation, or zero if
      // none needed.
      int32 precomputed_indexes_index;
  
      std::vector<Index> output_indexes;      // Indexes that this step outputs.
      std::vector<int32> output_cindex_ids;   // cindex_ids corresponding to each
                                              // of the output indexes.
  
      // If this component is of type kDescriptor (and note that the top-level
      // Descriptor is a concatenation over >= 1 parts), then we set value_parts
      // to a list of submatrix-indexes, each for the corresponding part of the
      // value.  If there is only one part, it will have one element which will be
      // the same as "value".
      std::vector<int32> value_parts;
      // deriv_parts is as "value_parts", but for parts of the derivative (if
      // we're doing backprop).
      std::vector<int32> deriv_parts;
  
      // for nodes corresponding to descriptors, input_locations_list will contain
      // information about the inputs to this descriptor, telling us for each row
      // of the matrix what other matrix rows it is a summation over.  this is a
      // quantity indexed[part-index][row-index], then a list of pairs (step,
      // row-index), representing source Cindexes present in a summation, that we
      // store here to avoid computing it twice in forward and backprop.
      std::vector<std::vector<std::vector<std::pair<int32,int32> > > > input_locations_list;
  
      StepInfo(): node_index(-1), value(0), deriv(0), segment(0),
                  precomputed_indexes_index(0) { }
    };
  
    // Computes the set of step-indexes of preceding steps that this step depends
    // on.  Assumes CreateLocationInfo() has already been called.  Requires
    // 'step_index' only to handle a special case, that if 'this_step' is a
    // component step, then the only step it depends on is the preceding step
    // (which is the component-input step).
    void ComputeStepDependencies(const std::vector<int32> &this_step,
                                 int32 step_index,
                                 unordered_set<int32> *dep_steps);
  
    // This function outputs to each element of "deriv_needed" a bool saying
    // whether, for that step, we need to allocate the matrix of derivatives
    // (interpret this as being at the output of that step).  This variable
    // also tells us whether we need to execute the backprop code for that step.
    //  'steps' is a vector of steps; each step is a list of cindexes.
    //  'step_to_segment', which should have the same dimension as 'steps',
    //    maps from step index to the segment it occurs in (only interesting
    //    for multi-segment/online computations).
    //  'deriv_needed' will be given the same length as 'steps'.
    void ComputeDerivNeeded(const std::vector<std::vector<int32> > &steps,
                            const std::vector<int32> &step_to_segment,
                            std::vector<bool> *deriv_needed);
  
    // this sets up steps_, destroying the input "by_step" in the process.  It
    // also sets various matrix and sub-matrix sizes in "computation".  The input
    // 'by_step' is elsewhere referred to as just 'step'; it is a vector of steps,
    // and each step is a vector of cindex_ids that are computed by that step.
    void CreateStepInfo(const std::vector<bool> &deriv_needed,
                        const std::vector<int32> &step_to_segment,
                        std::vector<std::vector<int32> > *by_step,
                        NnetComputation *computation);
  
    // Gets the stride type, kDefaultStride or kStrideEqualNumCols,
    // at the output of this node: interrogates component flags
    // looking for kInputContiguous or kOutputContiguous.
    MatrixStrideType GetStrideType(int32 node_index) const;
  
  
    // Miscellaneous info pertaining to various steps of the computation.  Indexed
    // by step-index.
    std::vector<StepInfo> steps_;
  
    /// This maps each cindex_id to its location.  However, you should not rely on
    /// its accuracy for cindex_ids that correspond to the Descriptors at
    /// Component inputs, since it's possible in principle for such cindex_ids to
    /// exist at >1 location.  (This is not a problem in practice, because we only
    /// need this for the outputs of component-nodes, and for computation inputs).
    /// A location is a pair (step-index, matrix-row-index).
    std::vector<std::pair<int32, int32> > cindex_id_to_location_;
  
  
    // Adds to the computation object the information about the matrix sizes
    void DefineMatrices(NnetComputation *computation) const;
  
    // Sets up sub-matrix indexes for nodes of type Descriptor (needed mainly
    // because Descriptors in general have many parts corresponding to
    // feature-dimension ranges, and they live in sub-matrices.
    void DefineSubmatrices(NnetComputation *computation);
  
    // Adds to the computation object the commands to allocate the matrices.
    // 'whole_submatrices' is as created by computation->GetWholeSubmatrices(), it
    // gives us the index of a submatrix containing the whole of each matrix.
    void AllocateMatrices(const std::vector<int32> &whole_submatrices,
                          NnetComputation *computation) const;
  
    // Sets up the precomputed indexes for each component, and sets the
    // precomputed_indexes_index value for each step.
    void SetUpPrecomputedIndexes(const std::vector<int32> &step_to_segment,
                                 NnetComputation *computation);
  
    // Adds to "computation" the command(s) for the forward computation
    // for this step.
    void CompileForward(int32 step, NnetComputation *computation) const;
  
    // Called from CompileForward, handles the case where the step corresponds
    // to a Component.
    void AddForwardStepComponent(int32 step, NnetComputation *computation) const;
  
    // Called from CompileForward, handles the case where the step corresponds
    // to an input node.
    void AddForwardStepInput(int32 step, NnetComputation *computation) const;
  
    // Returns true if step 'step' is an input step.   If step >= steps_.size(),
    // returns false.
    bool IsInputStep(int32 step) const;
  
  
    // Called from CompileForward, handles the case where the step
    // corresponds to type kDescriptor
    void CompileForwardDescriptor(
        int32 step, NnetComputation *computation) const;
  
    void CompileForwardSumDescriptor(
        int32 step, int32 part_index, NnetComputation *computation) const;
  
  
    // For the "part_index"'th part of the Descriptor for step "step" (which
    // must correspond to a Descriptor and not an Input or Component), this
    // function computes a vector of lists of submatrix locations of the inputs.
    // It is indexed by the number of rows in the output of this descriptor,
    // and the i'th element of the output is a list of pairs (step-index,
    // row-index-of-matrix).  The output of this row of this row of this part
    // of the computation will be a sum over those pairs.
    void ComputeInputLocationsList(
        int32 step, int32 part_index,
        std::vector<std::vector<std::pair<int32, int32> > > *input_locations)
        const;
  
    /**
       This function helps to handle scalar factors in Descriptors (expressions
       like `Scale(-1.0, <descriptor)`).  It splits an input_locations_list
       for one SumDescriptor (consisting of one of the appended-together parts
       of a Descriptor) by scale, such that each split-up locations_list
       corresponds to a single scaling factor.
       The scaling factors are all 1.0 by default, but may be different from
       1.0 if the user uses `Scale(...)` expressions in descriptors, e.g.
       `Scale(-1.0, lstm1.z)`.
       To help efficiency, this function treats the case where all the scales
       in the expression are the same (usually 1.0), as a special case.  In this
       case, 'split_locations_lists' will be empty and the shared scale (e.g. 1.0)
       is returned.
  
         @param [in] descriptor  The SumDescriptor for which we're getting
                        scalar factors.
         @param [in] input_locations_list This is one element of the
                        input_locations_list from the StepInfo of the step we are
                        computing, corresponding to this SumDescriptor (i.e. one
                        part of the Descriptor).  It is indexed by row-index, then
                        it is a list of pairs (step, row-index), representing
                        source Cindexes of a summation.  This function will work
                        out what scaling factors the pairs in these lists have.
         @param [out] split_locations_lists
                        We write to this location.  If all the scaling factors
                        are the same this will be set to the empty list and the
                        common scaling factor returned.  Otherwise +infinity
                        will be returned and the split-up list will be
                        written to the location.  Each element
                        (*split_locations_lists)[i] will be set to a pair
                        (alpha, partial_input_locations_list)
                        where alpha is the scaling factor associated with this
                        split-up piece (e.g. -1.0 if it was part of an expression
                        like `Scale(-1.0, lstm1.z)`), and
                        'partial_input_locations_list' is a vector with the same
                        dimension as 'input_locations_list' (indexed by row-index),
                        where partial_input_locations_list[r] will contain a subset
                        of the pairs present in input_locations_list[r], and
                        if we were to append together all the
                        (*split_locations_lists)[*].second.partial_input_locations_list[r],
                        we'd get a list with the same members as
                        input_locations_list[r], although not necessarily in the same
                        order.
          @return  In the general case (where multiple scales are used), returns
                   +infinity and sets 'split_locations_lists' to the split-up list.
                   In the special, but more common case where only a single scale
                   is used, return that scale (1.0 will be the most common value)
                   and set 'split_locations_lists' to empty; in this special case,
                   which has been made a special case for efficiency reasons,
                   the user should directly use the un-split locations list in
                   'input_locations_list'.
     */
    BaseFloat SplitByScale(const SumDescriptor &descriptor,
     const std::vector<std::vector<std::pair<int32,int32> > > &input_locations_list,
     std::vector<std::pair<BaseFloat,
                           std::vector<std::vector<std::pair<int32,int32> > > > >
                           *split_locations_lists) const;
  
    // Changes the format of the location-list produced by ComputeInputLocationsList,
    // to have pairs (sub-matrix, row) instead of (step, row), by replacing each step
    // (i.e. the first of each pair) with steps_[step].value.
    void ComputeValueSubmatLocationsList(
   const std::vector<std::vector<std::pair<int32, int32> > > &input_locations_list,
       std::vector<std::vector<std::pair<int32, int32> > > *submat_locations_list)
        const;
  
  
    // Changes the format of the location-list produced by
    // ComputeInputLocationsList, to have pairs (sub-matrix, row) instead of
    // (step, row), but with locations of derivatives not values (for use in
    // backprop).  It does this by replacing each step (i.e. the first of each
    // pair) with steps_[step].deriv, but if this value is zero (i.e. no such
    // derivative exists) it removes the pair.  This could occur in situations
    // where we only need to propagate the derivative selectively to some inputs.
    void ComputeDerivSubmatLocationsList(
   const std::vector<std::vector<std::pair<int32, int32> > > &input_locations_list,
   std::vector<std::vector<std::pair<int32, int32> > > *submat_locations_list)
        const;
  
  
  
    /** Adds to 'computation' commands for part of the forward computation
        corresponding to a Descriptor.  This is called from
        CompileForwardSumDescriptor.
  
        @param [in] value_submatrix_index  The submatrix index
                 of the quanitity we are computing (part of a Descriptor;
                 it's something like Sum(tdnn1, tdnn2) in general).
        @param [in] alpha  The scale (1.0 unless Scale(...) expressions are
                 involved in descriptors) with which these terms are present
                 in the summation.
        @param [in] submat_locations  Indexed by the row index of
                 the submatrix referred to by 'value_submatrix_index', each element is
                 a list of sources over which we must sum to obtain
                 that row.  Each source is a pair (submatrix-index, row-index).
    */
    void CompileForwardFromSubmatLocationsList(
        int32 value_submatrix_index,
        BaseFloat alpha,
        const std::vector<std::vector<std::pair<int32, int32> > > &submat_locations,
        NnetComputation *computation) const;
  
    /** Adds to 'computation' commands for part of the forward computation
        corresponding to a Descriptor.  This is called from
        CompileForwardFromSubmatLocationsList.
  
        @param [in] value_submatrix_index  The submatrix index
                 of the quanitity we are computing (part of a Descriptor;
                 it's something like Sum(tdnn1, tdnn2) in general).
        @param [in] alpha  The scale (1.0 unless Scale(...) expressions are
                 involved in descriptors) with which these terms are present
                 in the summation.
        @param [in] submat_locations  Indexed by the row index corresponding
                 to the rows of the submatrix referred to by 'value_submatrix_index',
                 this reprenents the source vector which we are adding to this row,
                 in the format (submatrix-index, row-index), or (-1, -1)
                 if in this case there is nothing to add.
         @param [in,out] computation  The computation which we are adding
                 commands to.
    */
    void CompileForwardFromSubmatLocations(
        int32 value_submatrix_index,
        BaseFloat alpha,
        const std::vector<std::pair<int32, int32> > &submat_locations,
        NnetComputation *computation) const;
  
  
    /** Adds to `computation` a command that adds to the submatrix in
        `value_submatrix_index` a quantity consisting of alpha times
        the submatrix in `input_submatrix_index`, with a row mapping
        given by `indexes`.
    */
    void CompileForwardFromIndexes(
        int32 value_submatrix_index,
        int32 input_submatrix_index,
        BaseFloat alpha,
        const std::vector<int32> &indexes,
        NnetComputation *computation) const;
  
  
    // Adds to "computation" the command(s) for the backward computation (if any) for
    // this step.  (non-const only because we clear the cached submat_locations).
    void CompileBackward(int32 step, NnetComputation *computation);
  
    // Called from CompileBackward, handles the case where the step corresponds
    // to a Component.
    void AddBackwardStepComponent(int32 step, NnetComputation *computation) const;
  
    // Called from CompileBackward, handles the case where the step
    // corresponds to an input.  If applicable, this generates a command for the
    // network to provide the derivative w.r.t. the input, to the user.
    void AddBackwardStepInput(int32 step, NnetComputation *computation) const;
  
    // Called from CompileBackward, handles the case where the step
    // corresponds to type kDescriptor.
    void CompileBackwardDescriptor(
        int32 step, NnetComputation *computation);
  
    // Called from CompileBackwardSumDescriptor.
    void CompileBackwardSumDescriptor(
        int32 step, int32 part_index,
        NnetComputation *computation) const;
  
    // Called from CompileBackwardForwardingDescriptor.
    void CompileBackwardFromSubmatLocationsList(
        int32 deriv_submatrix_index,
        BaseFloat alpha,
        const std::vector<std::vector<std::pair<int32, int32> > >&submat_locations,
        NnetComputation *computation) const;
  
  
    void CompileBackwardFromSubmatLocations(
        int32 deriv_submatrix_index,
        BaseFloat alpha,
        const std::vector<std::pair<int32, int32> > &submat_locations,
        NnetComputation *computation) const;
  
    // Called from CompileBackwardFromSubmatLocations - special case where
    // input is from just one matrix.
    void CompileBackwardFromIndexes(
        int32 deriv_submatrix_index,
        int32 input_deriv_submatrix_index,
        BaseFloat alpha,
        const std::vector<int32> &indexes,
        NnetComputation *computation) const;
  
  
    // [to be called after steps_ is set up and all the forward and backprop
    // commands have been added].  Adds to the computation the commands that
    // deinitialize all the matrices, except those that may be requested by
    // the user after the computation is done (i.e. outputs of the network,
    // and input derivatives).
    // 'whole_submatrices' is as created by computation->GetWholeSubmatrices(), it
    // gives us the index of a submatrix containing the whole of each matrix.
    void DeallocateMatrices(const std::vector<int32> &whole_submatrices,
                            const std::vector<int32> &step_to_segment,
                            NnetComputation *computation);
  
    // sets up the debug_info member of "computation".
    void OutputDebugInfo(NnetComputation *computation) const;
  
    void AddCommands(const std::vector<bool> &deriv_needed,
                     const std::vector<int32> &step_to_segment,
                     NnetComputation *computation);
  
  };
  
  
  
  
  } // namespace nnet3
  } // namespace kaldi
  
  
  #endif