// cudadecoder/cuda-fst.cc
//
// Copyright (c) 2019, NVIDIA CORPORATION.  All rights reserved.
// Hugo Braun, Justin Luitjens, Ryan Leary
//
// 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
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#if HAVE_CUDA == 1

#include "cudadecoder/cuda-fst.h"

#include <cuda_runtime_api.h>
#include <nvToolsExt.h>

namespace kaldi {
namespace cuda_decoder {

void CudaFst::ComputeOffsets(const fst::Fst<StdArc> &fst) {
  // count states since Fst doesn't provide this functionality
  num_states_ = 0;
  for (fst::StateIterator<fst::Fst<StdArc> > iter(fst); !iter.Done();
       iter.Next())
    ++num_states_;

  // allocate and initialize offset arrays
  h_final_.resize(num_states_);
  h_e_offsets_.resize(num_states_ + 1);
  h_ne_offsets_.resize(num_states_ + 1);

  // iterate through states and arcs and count number of arcs per state
  e_count_ = 0;
  ne_count_ = 0;

  // Init first offsets
  h_ne_offsets_[0] = 0;
  h_e_offsets_[0] = 0;
  for (int i = 0; i < num_states_; i++) {
    h_final_[i] = fst.Final(i).Value();
    // count emiting and non_emitting arcs
    for (fst::ArcIterator<fst::Fst<StdArc> > aiter(fst, i); !aiter.Done();
         aiter.Next()) {
      StdArc arc = aiter.Value();
      int32 ilabel = arc.ilabel;
      if (ilabel != 0) {  // emitting
        e_count_++;
      } else {  // non-emitting
        ne_count_++;
      }
    }
    h_ne_offsets_[i + 1] = ne_count_;
    h_e_offsets_[i + 1] = e_count_;
  }

  // We put the emitting arcs before the nonemitting arcs in the arc list
  // adding offset to the non emitting arcs
  // we go to num_states_+1 to take into account the last offset
  for (int i = 0; i < num_states_ + 1; i++)
    h_ne_offsets_[i] += e_count_;  // e_arcs before

  arc_count_ = e_count_ + ne_count_;
}

void CudaFst::AllocateData(const fst::Fst<StdArc> &fst) {
  d_e_offsets_ = static_cast<unsigned int *>(CuDevice::Instantiate().Malloc(
      (num_states_ + 1) * sizeof(*d_e_offsets_)));
  d_ne_offsets_ = static_cast<unsigned int *>(CuDevice::Instantiate().Malloc(
      (num_states_ + 1) * sizeof(*d_ne_offsets_)));
  d_final_ = static_cast<float *>(
      CuDevice::Instantiate().Malloc((num_states_) * sizeof(*d_final_)));

  h_arc_weights_.resize(arc_count_);
  h_arc_nextstate_.resize(arc_count_);
  // ilabels (id indexing)
  h_arc_id_ilabels_.resize(arc_count_);
  h_arc_olabels_.resize(arc_count_);

  d_arc_weights_ = static_cast<float *>(
      CuDevice::Instantiate().Malloc(arc_count_ * sizeof(*d_arc_weights_)));
  d_arc_nextstates_ = static_cast<StateId *>(
      CuDevice::Instantiate().Malloc(arc_count_ * sizeof(*d_arc_nextstates_)));

  // Only the ilabels for the e_arc are needed on the device
  d_arc_pdf_ilabels_ = static_cast<int32 *>(
      CuDevice::Instantiate().Malloc(e_count_ * sizeof(*d_arc_pdf_ilabels_)));
}

void CudaFst::PopulateArcs(const fst::Fst<StdArc> &fst) {
  // now populate arc data
  int e_idx = 0;
  int ne_idx = e_count_;  // starts where e_offsets_ ends
  for (int i = 0; i < num_states_; i++) {
    for (fst::ArcIterator<fst::Fst<StdArc> > aiter(fst, i); !aiter.Done();
         aiter.Next()) {
      StdArc arc = aiter.Value();
      int idx;
      if (arc.ilabel != 0) {  // emitting
        idx = e_idx++;
      } else {
        idx = ne_idx++;
      }
      h_arc_weights_[idx] = arc.weight.Value();
      h_arc_nextstate_[idx] = arc.nextstate;
      h_arc_id_ilabels_[idx] = arc.ilabel;
      // For now we consider id indexing == pdf indexing
      // If the two are differents, we'll call ApplyTransModelOnIlabels with a
      // TransitionModel
      h_arc_pdf_ilabels_[idx] = arc.ilabel;
      h_arc_olabels_[idx] = arc.olabel;
    }
  }
}

void CudaFst::ApplyTransitionModelOnIlabels(
    const TransitionModel &trans_model) {
  // Converting ilabel here, to avoid reindexing when reading nnet3 output
  // We only need to convert the emitting arcs
  // The emitting arcs are the first e_count_ arcs
  for (int iarc = 0; iarc < e_count_; ++iarc)
    h_arc_pdf_ilabels_[iarc] =
        trans_model.TransitionIdToPdf(h_arc_id_ilabels_[iarc]);
}

void CudaFst::CopyDataToDevice() {
  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMemcpy(
      d_e_offsets_, &h_e_offsets_[0], (num_states_ + 1) * sizeof(*d_e_offsets_),
      cudaMemcpyHostToDevice));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMemcpy(
      d_ne_offsets_, &h_ne_offsets_[0],
      (num_states_ + 1) * sizeof(*d_ne_offsets_), cudaMemcpyHostToDevice));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMemcpy(d_final_, &h_final_[0],
                                                num_states_ * sizeof(*d_final_),
                                                cudaMemcpyHostToDevice));

  KALDI_DECODER_CUDA_API_CHECK_ERROR(
      cudaMemcpy(d_arc_weights_, &h_arc_weights_[0],
                 arc_count_ * sizeof(*d_arc_weights_), cudaMemcpyHostToDevice));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMemcpy(
      d_arc_nextstates_, &h_arc_nextstate_[0],
      arc_count_ * sizeof(*d_arc_nextstates_), cudaMemcpyHostToDevice));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMemcpy(
      d_arc_pdf_ilabels_, &h_arc_pdf_ilabels_[0],
      e_count_ * sizeof(*d_arc_pdf_ilabels_), cudaMemcpyHostToDevice));
}

void CudaFst::Initialize(const fst::Fst<StdArc> &fst,
                         const TransitionModel *trans_model) {
  nvtxRangePushA("CudaFst constructor");
  start_ = fst.Start();

  ComputeOffsets(fst);
  AllocateData(fst);
  // Temporarily allocating data for this vector
  // We just need it during CSR generation. We will clear it
  // at the end of Initialize
  h_arc_pdf_ilabels_.resize(arc_count_);
  PopulateArcs(fst);
  if (trans_model) ApplyTransitionModelOnIlabels(*trans_model);

  KALDI_ASSERT(d_e_offsets_);
  KALDI_ASSERT(d_ne_offsets_);
  KALDI_ASSERT(d_final_);
  KALDI_ASSERT(d_arc_weights_);
  KALDI_ASSERT(d_arc_nextstates_);
  KALDI_ASSERT(d_arc_pdf_ilabels_);

  CopyDataToDevice();

  // Making sure the graph is ready
  cudaDeviceSynchronize();
  KALDI_DECODER_CUDA_CHECK_ERROR();
  h_arc_pdf_ilabels_.clear();  // we don't need those on host
  nvtxRangePop();
}

void CudaFst::Finalize() {
  nvtxRangePushA("CudaFst destructor");

  // Making sure that Initialize was called before Finalize
  KALDI_ASSERT(d_e_offsets_ &&
               "Please call CudaFst::Initialize() before calling Finalize()");
  KALDI_ASSERT(d_ne_offsets_);
  KALDI_ASSERT(d_final_);
  KALDI_ASSERT(d_arc_weights_);
  KALDI_ASSERT(d_arc_nextstates_);
  KALDI_ASSERT(d_arc_pdf_ilabels_);

  CuDevice::Instantiate().Free(d_e_offsets_);
  CuDevice::Instantiate().Free(d_ne_offsets_);
  CuDevice::Instantiate().Free(d_final_);
  CuDevice::Instantiate().Free(d_arc_weights_);
  CuDevice::Instantiate().Free(d_arc_nextstates_);
  CuDevice::Instantiate().Free(d_arc_pdf_ilabels_);
  nvtxRangePop();
}

}  // end namespace cuda_decoder
}  // end namespace kaldi

#endif  // HAVE_CUDA == 1