// chain/chain-denominator.cc
  
  // Copyright      2015   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.
  
  
  #include "chain/chain-denominator.h"
  #include "chain/chain-kernels-ansi.h"
  
  namespace kaldi {
  namespace chain {
  
  
  DenominatorComputation::DenominatorComputation(
      const ChainTrainingOptions &opts,
      const DenominatorGraph &den_graph,
      int32 num_sequences,
      const CuMatrixBase<BaseFloat> &nnet_output):
      opts_(opts),
      den_graph_(den_graph),
      num_sequences_(num_sequences),
      frames_per_sequence_(nnet_output.NumRows() / num_sequences_),
      nnet_output_deriv_transposed_(
          nnet_output.NumCols(),
          std::min<int32>(nnet_output.NumRows(),
                          static_cast<int32>(kMaxDerivTimeSteps) *
                          num_sequences_)),
      alpha_(frames_per_sequence_ + 1,
             den_graph_.NumStates() * num_sequences_ + num_sequences_,
             kUndefined),
      beta_(2, den_graph_.NumStates() * num_sequences_ + num_sequences_,
            kUndefined),
      tot_prob_(num_sequences_, kUndefined),
      tot_log_prob_(num_sequences_, kUndefined),
      log_correction_term_(num_sequences_, kUndefined),
      ok_(true) {
    // We don't let leaky_hmm_coefficient be exactly zero (although that would
    // make sense mathematically, corresponding to "turning off" the leaky HMM),
    // because that would lead to underflow and eventually NaN's or inf's
    // appearing in the computation, since we do this computation not in
    // log-space.
    KALDI_ASSERT(opts_.leaky_hmm_coefficient > 0.0 &&
                 opts_.leaky_hmm_coefficient < 1.0);
  
    if (RandInt(0, 99) == 0) {
      // A check, that all values in nnet_output are in the range [-30, 30]..
      // otherwise derivatives will be wrong (search below for 30).
      BaseFloat max_val = nnet_output.Max(), min_val = nnet_output.Min();
      if (max_val > 30.0 || min_val < -30.0) {
        KALDI_WARN << "Nnet outputs " << min_val << ", "
                   << max_val <<
            " outside the range [-30,30], derivs may be inaccurate.";
      }
    }
  
    // make sure the alpha sums and beta sums are zeroed.
    alpha_.ColRange(den_graph_.NumStates() * num_sequences_,
                    num_sequences_).SetZero();
    beta_.ColRange(den_graph_.NumStates() * num_sequences_,
                   num_sequences_).SetZero();
  
    KALDI_ASSERT(nnet_output.NumRows() % num_sequences == 0);
    // the kStrideEqualNumCols argument is so that we can share the same
    // memory block with xent_output_deriv (see chain-training.cc, search for
    // kStrideEqualNumCols).  This depends on how the allocator works, and
    // actually might not happen, but anyway, the impact on speed would
    // likely be un-measurably small.
    exp_nnet_output_transposed_.Resize(nnet_output.NumCols(),
                                       nnet_output.NumRows(),
                                       kUndefined, kStrideEqualNumCols);
    exp_nnet_output_transposed_.CopyFromMat(nnet_output, kTrans);
    // We limit the nnet output to the range [-30,30] before doing the exp;
    // this avoids NaNs appearing in the forward-backward computation, which
    // is not done in log space.
    exp_nnet_output_transposed_.ApplyExpLimited(-30.0, 30.0);
  }
  
  
  void DenominatorComputation::AlphaFirstFrame() {
    // dim == num_hmm_states_ * num_sequences_.
    BaseFloat *first_frame_alpha = alpha_.RowData(0);
    // create a 'fake matrix' - view this row as a matrix.
    // initializer takes [pointer, num-rows, num-cols, stride].
    CuSubMatrix<BaseFloat> alpha_mat(first_frame_alpha,
                                     den_graph_.NumStates(),
                                     num_sequences_,
                                     num_sequences_);
    // TODO (possible): It would be more efficient here if we implemented a
    // CopyColsFromVec function in class CuMatrix.
    alpha_mat.SetZero();
    alpha_mat.AddVecToCols(1.0, den_graph_.InitialProbs(), 0.0);
  }
  
  
  // the alpha computation for some 0 < t <= num_time_steps_.
  void DenominatorComputation::AlphaGeneralFrame(int32 t) {
    KALDI_ASSERT(t > 0 && t <= frames_per_sequence_);
    BaseFloat *this_alpha = alpha_.RowData(t);
    const BaseFloat *prev_alpha_dash = alpha_.RowData(t - 1);
    const Int32Pair *backward_transitions = den_graph_.BackwardTransitions();
    const DenominatorGraphTransition *transitions = den_graph_.Transitions();
    int32 num_pdfs = exp_nnet_output_transposed_.NumRows(),
        num_hmm_states = den_graph_.NumStates(),
        num_sequences = num_sequences_;
  
    // 'probs' is the matrix of pseudo-likelihoods for frame t - 1.
    CuSubMatrix<BaseFloat> probs(exp_nnet_output_transposed_, 0, num_pdfs,
                                 (t-1) * num_sequences_, num_sequences_);
    const BaseFloat *prob_data = probs.Data();
  
  #if HAVE_CUDA == 1
    if (CuDevice::Instantiate().Enabled()) {
      CuTimer tim;
      dim3 dimBlock(std::min<int32>(CU1DBLOCK, num_sequences), 1, 1);
      dim3 dimGrid(n_blocks(num_sequences, dimBlock.x), num_hmm_states, 1);
  
      while (1) {
        if (dimGrid.y > 65535)  // the hardware doesn't allow more than this.
          dimGrid.y = 65535;
        cuda_chain_hmm_forward(dimGrid, dimBlock,
                               backward_transitions, transitions,
                               num_sequences, den_graph_.NumStates(),
                               prob_data, probs.Stride(), prev_alpha_dash,
                               this_alpha);
        CU_SAFE_CALL(cudaGetLastError());
        if (dimGrid.y == num_hmm_states) {
          break;  // this is the normal case.
        } else {
          // We reach this code only in the unusual case where num_hmm_states >
          // 65535.  We can compute the alphas for the remaining HMM states by
          // moving some of the array pointers and making the call again.
          backward_transitions += dimGrid.y;
          this_alpha += dimGrid.y * num_sequences;
          num_hmm_states -= dimGrid.y;
          dimGrid.y = num_hmm_states;
        }
      }
      CuDevice::Instantiate().AccuProfile(__func__, tim);
    } else
  #endif
    {
      int32 prob_stride = probs.Stride();
      for (int32 h = 0; h < num_hmm_states; h++) {
        for (int32 s = 0; s < num_sequences; s++) {
          double this_tot_alpha = 0.0;
          const DenominatorGraphTransition
              *trans_iter = transitions + backward_transitions[h].first,
              *trans_end = transitions + backward_transitions[h].second;
          for (; trans_iter != trans_end; ++trans_iter) {
            BaseFloat transition_prob = trans_iter->transition_prob;
            int32 pdf_id = trans_iter->pdf_id,
                prev_hmm_state = trans_iter->hmm_state;
            BaseFloat prob = prob_data[pdf_id * prob_stride + s],
                this_prev_alpha = prev_alpha_dash[prev_hmm_state * num_sequences + s];
            this_tot_alpha += this_prev_alpha * transition_prob * prob;
          }
          // Let arbitrary_scale be the inverse of the alpha-sum value that we
          // store in the same place we'd store the alpha for the state numbered
          // 'num_hmm_states'. We multiply this into all the
          // transition-probabilities from the previous frame to this frame, in
          // both the forward and backward passes, in order to keep the alphas in
          // a good numeric range.  This won't affect the posteriors, but when
          // computing the total likelihood we'll need to compensate for it later
          // on.
          BaseFloat arbitrary_scale =
              1.0 / prev_alpha_dash[num_hmm_states * num_sequences + s];
          KALDI_ASSERT(this_tot_alpha - this_tot_alpha == 0);
          this_alpha[h * num_sequences + s] = this_tot_alpha * arbitrary_scale;
        }
      }
    }
  }
  
  void DenominatorComputation::AlphaDash(int32 t) {
    BaseFloat *this_alpha = alpha_.RowData(t);
  
    // create a 'fake matrix' for the regular alphas- view this row as a matrix.
    // initializer takes [pointer, num-rows, num-cols, stride].
    CuSubMatrix<BaseFloat> alpha_mat(this_alpha,
                                     den_graph_.NumStates(),
                                     num_sequences_,
                                     num_sequences_);
  
    // the alpha-dash is the sum of alpha over all states.
    CuSubVector<BaseFloat> alpha_sum_vec(this_alpha +
                                         den_graph_.NumStates() * num_sequences_,
                                         num_sequences_);
    alpha_sum_vec.AddRowSumMat(1.0, alpha_mat, 0.0);
  
    alpha_mat.AddVecVec(opts_.leaky_hmm_coefficient,
                        den_graph_.InitialProbs(),
                        alpha_sum_vec);
    // it's now alpha-dash.
  }
  
  // compute beta from beta-dash.
  void DenominatorComputation::Beta(int32 t) {
    BaseFloat *this_beta_dash = beta_.RowData(t % 2);
    // create a 'fake matrix' for the regular beta-dash (which is
    // the counterpart of alpha-dash)- view this row as a matrix.
    // initializer takes [pointer, num-rows, num-cols, stride].
    CuSubMatrix<BaseFloat> beta_dash_mat(this_beta_dash,
                                         den_graph_.NumStates(),
                                         num_sequences_,
                                         num_sequences_);
    // making the t index implicit, the beta-dash-sum for each sequence is the sum
    // over all states i of beta_i * opts_.leaky_hmm_coefficient * initial_prob_i.
    CuSubVector<BaseFloat> beta_dash_sum_vec(
        this_beta_dash + den_graph_.NumStates() * num_sequences_,
        num_sequences_);
    beta_dash_sum_vec.AddMatVec(opts_.leaky_hmm_coefficient, beta_dash_mat,
                                kTrans, den_graph_.InitialProbs(), 0.0);
    // we are computing beta in place.  After the following, beta-dash-mat
    // will contain the actual beta (i.e. the counterpart of alpha),
    // not the beta-dash.
    beta_dash_mat.AddVecToRows(1.0, beta_dash_sum_vec);
  }
  
  BaseFloat DenominatorComputation::Forward() {
    AlphaFirstFrame();
    AlphaDash(0);
    for (int32 t = 1; t <= frames_per_sequence_; t++) {
      AlphaGeneralFrame(t);
      AlphaDash(t);
    }
    return ComputeTotLogLike();
  }
  
  BaseFloat DenominatorComputation::ComputeTotLogLike() {
    tot_prob_.Resize(num_sequences_);
    // View the last alpha-dash as a matrix of size num-hmm-states by num-sequences.
    CuSubMatrix<BaseFloat> last_alpha_dash(
        alpha_.RowData(frames_per_sequence_),
        den_graph_.NumStates(),
        num_sequences_,
        num_sequences_);
  
    tot_prob_.AddRowSumMat(1.0, last_alpha_dash, 0.0);
    // we should probably add an ApplyLog() function that takes a vector argument.
    tot_log_prob_ = tot_prob_;
    tot_log_prob_.ApplyLog();
    BaseFloat tot_log_prob = tot_log_prob_.Sum();
  
    // We now have to add something for the arbitrary scaling factor.  [note: the
    // purpose of the arbitrary scaling factors was to keep things in a good
    // floating-point range]
    // The inverses of all the tot-alpha quantities, for t = 0
    // ... frames_per_sequence_ - 1, were included as the 'arbitrary factors' in
    // the transition-probs, so we need to multiply them all together (not
    // inversed) and add them as a correction term to the total log-likes.
    // These tot-alpha quantities were stored in the same place that we would
    // have stored the HMM-state numbered 'num_hmm_states'.
    int32 num_hmm_states = den_graph_.NumStates();
    CuSubMatrix<BaseFloat> inv_arbitrary_scales(
        alpha_, 0, frames_per_sequence_,
        num_sequences_ * num_hmm_states, num_sequences_);
    CuMatrix<BaseFloat> log_inv_arbitrary_scales(
        inv_arbitrary_scales);
    log_inv_arbitrary_scales.ApplyLog();
    BaseFloat log_inv_arbitrary_scales_product =
        log_inv_arbitrary_scales.Sum();
    return tot_log_prob + log_inv_arbitrary_scales_product;
  }
  
  
  
  bool DenominatorComputation::Backward(
      BaseFloat deriv_weight,
      CuMatrixBase<BaseFloat> *nnet_output_deriv) {
    BetaDashLastFrame();
    Beta(frames_per_sequence_);
    for (int32 t = frames_per_sequence_ - 1; t >= 0; t--) {
      BetaDashGeneralFrame(t);
      if (GetVerboseLevel() >= 1 || t == 0)
        BetaGeneralFrameDebug(t);
      Beta(t);
      if (t % kMaxDerivTimeSteps == 0) {
        // commit the derivative stored in nnet_output_deriv_transposed_ by adding
        // its transpose to the appropriate sub-matrix of 'nnet_output_deriv'.
        int32 chunk_frames = std::min<int32>(static_cast<int32>(kMaxDerivTimeSteps),
                                             frames_per_sequence_ - t),
                  num_pdfs = exp_nnet_output_transposed_.NumRows();
        CuSubMatrix<BaseFloat> transposed_deriv_part(
            nnet_output_deriv_transposed_,
            0, num_pdfs,
            0, chunk_frames * num_sequences_);
        CuSubMatrix<BaseFloat> output_deriv_part(
            *nnet_output_deriv,
            t * num_sequences_, chunk_frames * num_sequences_,
            0, num_pdfs);
        output_deriv_part.AddMat(deriv_weight, transposed_deriv_part, kTrans);
        if (t != 0)
          transposed_deriv_part.SetZero();
      }
    }
  
    return ok_;
  }
  
  void DenominatorComputation::BetaDashLastFrame() {
    // sets up the beta-dash quantity on the last frame (frame ==
    // frames_per_sequence_).  Note that the betas we use here contain a
    // 1/(tot-prob) factor in order to simplify the backprop.
  
    int32 t = frames_per_sequence_;
    BaseFloat *last_frame_beta_dash = beta_.RowData(t % 2);
  
    // create a 'fake matrix' - view this row as a matrix.
    CuSubMatrix<BaseFloat> beta_dash_mat(last_frame_beta_dash,
                                         den_graph_.NumStates(),
                                         num_sequences_,
                                         num_sequences_);
    CuVector<BaseFloat> inv_tot_prob(tot_prob_);
    inv_tot_prob.InvertElements();
    // the beta values at the end of the file only vary with the sequence-index,
    // not with the HMM-index.  We treat all states as having a final-prob of one.
    beta_dash_mat.CopyRowsFromVec(inv_tot_prob);
  }
  
  void DenominatorComputation::BetaDashGeneralFrame(int32 t) {
    KALDI_ASSERT(t >= 0 && t < frames_per_sequence_);
    int32 num_pdfs = exp_nnet_output_transposed_.NumRows();
    // t_wrapped gives us the time-index we use when indexing
    // nnet_output_deriv_transposed_; to save memory we limit the size of the
    // matrix, storing only chunks of frames at a time, and we add it to the
    // non-transposed output whenever we finish a chunk.
    int32 t_wrapped = t % static_cast<int32>(kMaxDerivTimeSteps);
    const BaseFloat *this_alpha_dash = alpha_.RowData(t),
        *next_beta = beta_.RowData((t + 1) % 2);
    BaseFloat *this_beta_dash = beta_.RowData(t % 2);
    const Int32Pair *forward_transitions = den_graph_.ForwardTransitions();
    const DenominatorGraphTransition *transitions = den_graph_.Transitions();
    // 'probs' is the matrix of pseudo-likelihoods for frame t.
    CuSubMatrix<BaseFloat> probs(exp_nnet_output_transposed_, 0, num_pdfs,
                                 t * num_sequences_, num_sequences_),
        log_prob_deriv(nnet_output_deriv_transposed_, 0, num_pdfs,
                       t_wrapped * num_sequences_, num_sequences_);
  
    int32 num_hmm_states = den_graph_.NumStates(),
        num_sequences = num_sequences_;
  
  #if HAVE_CUDA == 1
    if (CuDevice::Instantiate().Enabled()) {
      CuTimer tim;
      dim3 dimBlock(std::min<int32>(CU1DBLOCK, num_sequences), 1, 1);
      dim3 dimGrid(n_blocks(num_sequences, dimBlock.x), num_hmm_states, 1);
      while (1) {
        if (dimGrid.y > 65535)  // the hardware doesn't allow more than this.
          dimGrid.y = 65535;
        cuda_chain_hmm_backward(dimGrid, dimBlock, forward_transitions, transitions,
                                num_sequences, num_hmm_states,
                                probs.Data(), probs.Stride(),
                                this_alpha_dash, next_beta, this_beta_dash,
                                log_prob_deriv.Data(), log_prob_deriv.Stride());
        CU_SAFE_CALL(cudaGetLastError());
        if (dimGrid.y == num_hmm_states) {
          break;  // this is the normal case.
        } else {
          // We reach this code only in the unusual case where num_hmm_states >
          // 65535.  We can compute the betas (and log-prob derivatives) for the
          // remaining HMM states by moving some of the array pointers and making
          // the call again.
          forward_transitions += dimGrid.y;
          this_alpha_dash += dimGrid.y * num_sequences;
          this_beta_dash += dimGrid.y * num_sequences;
          num_hmm_states -= dimGrid.y;
          dimGrid.y = num_hmm_states;
        }
      }
      CuDevice::Instantiate().AccuProfile(__func__, tim);
    } else
  #endif
    {
      int32 prob_stride = probs.Stride(),
           deriv_stride = log_prob_deriv.Stride();
      const BaseFloat *prob_data = probs.Data();
      BaseFloat *log_prob_deriv_data = log_prob_deriv.Data();
      for (int32 h = 0; h < num_hmm_states; h++) {
        for (int32 s = 0; s < num_sequences; s++) {
          BaseFloat this_alpha_dash_prob = this_alpha_dash[h * num_sequences + s],
              inv_arbitrary_scale =
              this_alpha_dash[num_hmm_states * num_sequences + s];
          double tot_variable_factor = 0.0;
          BaseFloat occupation_factor = this_alpha_dash_prob /
              inv_arbitrary_scale;
          const DenominatorGraphTransition
              *trans_iter = transitions + forward_transitions[h].first,
              *trans_end = transitions + forward_transitions[h].second;
          for (; trans_iter != trans_end; ++trans_iter) {
            BaseFloat transition_prob = trans_iter->transition_prob;
            int32 pdf_id = trans_iter->pdf_id,
                next_hmm_state = trans_iter->hmm_state;
            BaseFloat variable_factor = transition_prob *
                next_beta[next_hmm_state * num_sequences + s] *
                prob_data[pdf_id * prob_stride + s];
            tot_variable_factor += variable_factor;
            BaseFloat occupation_prob = variable_factor * occupation_factor;
            log_prob_deriv_data[pdf_id * deriv_stride + s] += occupation_prob;
          }
          this_beta_dash[h * num_sequences + s] =
              tot_variable_factor / inv_arbitrary_scale;
        }
      }
    }
  }
  
  void DenominatorComputation::BetaGeneralFrameDebug(int32 t) {
    BaseFloat num_hmm_states = den_graph_.NumStates(),
        alpha_beta_size = num_hmm_states * num_sequences_;
    CuSubVector<BaseFloat> this_alpha_dash(alpha_.RowData(t), alpha_beta_size),
        this_beta_dash(beta_.RowData(t % 2), alpha_beta_size);
    int32 t_wrapped = t % static_cast<int32>(kMaxDerivTimeSteps),
        num_pdfs = exp_nnet_output_transposed_.NumRows();
    CuSubMatrix<BaseFloat> this_log_prob_deriv(
        nnet_output_deriv_transposed_, 0, num_pdfs,
        t_wrapped * num_sequences_, num_sequences_);
    BaseFloat alpha_beta_product = VecVec(this_alpha_dash,
                                          this_beta_dash),
        this_log_prob_deriv_sum = this_log_prob_deriv.Sum();
    if (!ApproxEqual(alpha_beta_product, num_sequences_)) {
      KALDI_WARN << "On time " << t << ", alpha-beta product "
                 << alpha_beta_product << " != " << num_sequences_
                 << " alpha-dash-sum = " << this_alpha_dash.Sum()
                 << ", beta-dash-sum = " << this_beta_dash.Sum();
      if (fabs(alpha_beta_product - num_sequences_) > 2.0) {
        KALDI_WARN << "Excessive error detected, will abandon this minibatch";
        ok_ = false;
      }
    }
    // use higher tolerance, since we are using randomized pruning for the
    // log-prob derivatives.
    if (!ApproxEqual(this_log_prob_deriv_sum,
                     num_sequences_, 0.01)) {
      KALDI_WARN << "On time " << t << ", log-prob-deriv sum "
                 << this_log_prob_deriv_sum << " != " << num_sequences_;
      if (fabs(this_log_prob_deriv_sum - num_sequences_) > 2.0) {
        KALDI_WARN << "Excessive error detected, will abandon this minibatch";
        ok_ = false;
      }
    }
  }
  
  
  }  // namespace chain
  }  // namespace kaldi