/*
 * Copyright (c) 2020-2023, NVIDIA CORPORATION.  All rights reserved.
 *
 * 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.
 */
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#pragma once

#include "moe_gemm_kernels.h"
#include <cuda_runtime_api.h>

#include "contrib_ops/cuda/bert/transformer_cuda_common.h"
#include "core/common/common.h"

#include "cutlass/numeric_types.h"

using namespace onnxruntime;

namespace ort_fastertransformer {

static inline size_t pad_to_multiple_of_16(size_t input) {
  static constexpr int ALIGNMENT = 16;
  return ALIGNMENT * ((input + ALIGNMENT - 1) / ALIGNMENT);
}

/*
  Launches the topk gating softmax required for the MoE layers.

  Params:
  input - a [num_rows x num_experts]
  finished - [num_rows] vector with 1 if the sentence at this row is done translating and 0 otherwise.
  output - a buffer of shape [num_rows x k] containing the top-k values of the softmax for each row.
  indices - a matrix of shape [num_rows x k] containing the top-k experts each row should get routed to.
  source_rows - a matrix of shape [num_rows x k] used internally for permuting. source_rows[row][k] =  k * num_rows +
  row. It is constructed like this so we can track where each of the original rows end up in order to perform the
                "k-way" reduction later in the routing.

  num_rows - The number of rows in the matrix
  num_experts - The number of expert layers present
  k - k value in topk
*/
template <typename T>
void topk_gating_softmax_kernelLauncher(const T* input, const bool* finished, T* output, T* softmax_temp_out,
                                        int* indices, int* source_row, int num_rows, int num_experts, int k,
                                        cudaStream_t stream);

class CubKeyValueSorter {
 public:
  CubKeyValueSorter();

  CubKeyValueSorter(int num_experts);

  void update_num_experts(int num_experts);

  size_t getWorkspaceSize(const size_t num_key_value_pairs);

  void run(void* workspace, const size_t workspace_size, const int* keys_in, int* keys_out, const int* values_in,
           int* values_out, const size_t num_key_value_pairs, cudaStream_t stream);

 private:
  size_t num_key_value_pairs_;
  int num_experts_;
  int num_bits_;
};

template <typename T>
void initialize_moe_routing_kernelLauncher(const T* unpermuted_input, T* permuted_output,
                                           const int* expanded_dest_row_to_expanded_source_row,
                                           int* expanded_source_row_to_expanded_dest_row, int num_rows, int active_rows,
                                           int cols, int k, cudaStream_t stream);

template <typename T>
void finalize_moe_routing_kernelLauncher(const T* expanded_permuted_rows, T* reduced_unpermuted_output, const T* bias,
                                         const T* scales, const int* expanded_source_row_to_expanded_dest_row,
                                         const int* expert_for_source_row, int num_rows, int cols, int k,
                                         cudaStream_t stream);

template <typename T>
void finalize_moe_routing_kernelLauncher(const T* expanded_permuted_rows, T* reduced_unpermuted_output, const T* skip,
                                         const T* bias, const T* scales,
                                         const int* expanded_source_row_to_expanded_dest_row,
                                         const int* expert_for_source_row, int num_rows, int cols, int k,
                                         cudaStream_t stream);

template <typename T>
void finalize_moe_routing_kernelLauncher(const T* expanded_permuted_rows, T* reduced_unpermuted_output, const T* skip_1,
                                         const T* skip_2, const T* bias, const T* scales,
                                         const int* expanded_source_row_to_expanded_dest_row,
                                         const int* expert_for_source_row, int num_rows, int cols, int k,
                                         cudaStream_t stream);

// Assumes inputs activations are row major. Weights need to be preprocessed by th_op/weight_quantize.cc .
// Nested in a class to avoid multiple calls to cudaGetDeviceProperties as this call can be expensive.
// Avoid making several duplicates of this class.
template <typename T,          /*The type used for activations/scales/compute*/
          typename WeightType, /* The type for the MoE weights */
          typename Enable = void>
class CutlassMoeFCRunner {
 public:
  CutlassMoeFCRunner(int sm_version, bool has_fc3, bool normalize_routing_weights, bool use_sparse_mixer);

  size_t getWorkspaceSize(size_t num_rows, size_t hidden_size, size_t inter_size, size_t num_experts, size_t k);

  void run_moe_fc(const T* input_activations, const T* gating_output, const WeightType* fc1_expert_weights,
                  const T* fc1_scales, const T* fc1_expert_biases, ActivationType fc1_activation_type,
                  const WeightType* fc3_expert_weights, const T* fc3_scales, const T* fc3_expert_biases,
                  const WeightType* fc2_expert_weights, const T* fc2_scales, int num_rows, int hidden_size,
                  int inter_size, int num_experts, int local_num_experts, int local_experts_start_index, int k,
                  char* workspace_ptr, T* fc2_result, T* expert_scales, int* expanded_source_row_to_expanded_dest_row,
                  int* expert_for_source_row, cudaStream_t stream);

  void run_moe_fc(const T* input_activations, const T* gating_output, const WeightType* fc1_expert_weights,
                  const T* fc1_scales, const T* fc1_expert_biases, ActivationType fc1_activation_type,
                  const WeightType* fc3_expert_weights, const T* fc3_scales, const T* fc3_expert_biases,
                  const WeightType* fc2_expert_weights, const T* fc2_scales, int num_rows, int hidden_size,
                  int inter_size, int num_experts, int local_num_experts, int local_experts_start_index, int k,
                  char* workspace_ptr, T* fc2_result, const bool* finished, int active_rows, T* expert_scales,
                  int* expanded_source_row_to_expanded_dest_row, int* expert_for_source_row, cudaStream_t stream);

  void compute_total_rows_before_expert(const int* sorted_indices, int total_indices, int num_experts,
                                        int64_t* total_rows_before_expert, cudaStream_t stream);

  void dispatch_activations(int64_t* total_rows_before_expert, int num_experts, int local_num_experts,
                            int local_experts_start_index, cudaStream_t stream);

  void get_total_rows_info(int64_t experts_start_index, int64_t local_num_experts, int64_t& total_past_rows,
                           int64_t& total_covered_rows);

 private:
  void configure_ws_ptrs(char* ws_ptr, size_t num_rows, size_t hidden_size, size_t inter_size, size_t num_experts,
                         size_t k);

 private:
  CubKeyValueSorter sorter_;
  MoeGemmRunner<T, WeightType> moe_gemm_runner_;

  // Pointers
  int* source_rows_;
  int* permuted_rows_;
  int* permuted_experts_;
  char* sorter_ws_;
  T* permuted_data_;
  T* softmax_out_;

  int64_t* total_rows_before_expert_;

  T* fc1_result_;
  T* fc3_result_;

  bool has_fc3_;
  bool normalize_routing_weights_;
  bool use_sparse_mixer_;

  // Cuda events
  contrib::cuda::AutoDestoryCudaEvent cuda_event_;

  int64_t total_past_rows_;
  int64_t total_covered_rows_;

  // TODO: use pinned memory
  std::vector<int64_t> total_rows_before_expert_host_;
};

template <typename WeightType>
class CutlassMoeFCRunner<float, WeightType, typename std::enable_if_t<!std::is_same<float, WeightType>::value>> {
 public:
  CutlassMoeFCRunner(int sm_version, bool has_fc3, bool normalize_routing_weights, bool use_sparse_mixer);

  size_t getWorkspaceSize(size_t num_rows, size_t hidden_size, size_t inter_size, size_t num_experts, size_t k) {
    return 0;
  }
};

}  // namespace ort_fastertransformer
