/*
 * Copyright (c) 2023, NVIDIA CORPORATION.
 *
 * 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.
 */

#include <cuda_runtime.h>

#include <core/hctr_impl/hctr_backend.hpp>
#include <embedding/common.hpp>
#include <utils.hpp>

namespace HugeCTR {
DataDistributionInput::DataDistributionInput(std::shared_ptr<core::CoreResourceManager> core,
                                             const std::vector<int> &dr_lookup_ids,
                                             core23::DataType key_type,
                                             core23::DataType offset_type)
    : num_lookup_(dr_lookup_ids.size()),
      dr_lookup_ids_(dr_lookup_ids),
      key_type(key_type),
      offset_type(offset_type) {
  CudaDeviceContext ctx(core->get_device_id());

  core23::Device device(core23::DeviceType::GPU, core->get_device_id());
  core23::TensorParams params = core23::TensorParams().device(device);

  // 2x for both keys & bucket_range
  this->d_ptrs_ = core23::Tensor(
      params.shape({static_cast<int64_t>(num_lookup_ * 2)}).data_type(core23::ScalarType::Pointer));
  this->h_ptrs_ = core23::Tensor(core23::TensorParams()
                                     .device(core23::DeviceType::CPU)
                                     .shape({static_cast<int64_t>(num_lookup_ * 2)})
                                     .data_type(core23::ScalarType::Pointer));
}

void DataDistributionInput::copy_tensor_vec(const std::vector<core23::Tensor> &dp_keys,
                                            const std::vector<core23::Tensor> &dp_bucket_range,
                                            cudaStream_t stream) {
  int num_lookup = dp_keys.size();
  HCTR_CHECK(num_lookup <= num_lookup_);

  // concat both arrays so we only need to copy up one array of pointers
  for (size_t i = 0; i < num_lookup_; ++i) {
    int dr_lookup_id = dr_lookup_ids_[i];
    HCTR_CHECK(dr_lookup_id < num_lookup);

    h_ptrs_.data<void *>()[i] = dp_keys[dr_lookup_id].data();
    // ATTENTION:
    // 1. When dp_bucket_range is generated by data distributor, its size is equal to the length of
    // phyical num lookup, which is generated by separating column-wise sharding table, so we should
    // use i to get dp_bucket_range.
    // 2. But when we change the data distributor to take dp_bucket_range from data reader, its
    // size is equal to the length of user configured num lookup. Then we should use dr_lookup_id to
    // index dp_bucket_range
    if (dp_bucket_range.size() == dp_keys.size()) {
      h_ptrs_.data<void *>()[num_lookup_ + i] = dp_bucket_range[dr_lookup_id].data();
    } else {
      HCTR_CHECK(num_lookup_ == dp_bucket_range.size());
      h_ptrs_.data<void *>()[num_lookup_ + i] = dp_bucket_range[i].data();
    }
  }

  core23::copy_async(d_ptrs_, h_ptrs_, stream);
}
}  // namespace HugeCTR

namespace embedding {

std::ostream &operator<<(std::ostream &os, const CompressionStrategy &p) {
  switch (p) {
    case CompressionStrategy::Reduction:
      os << "reduction";
      break;
    case CompressionStrategy::Unique:
      os << "unique";
      break;
    default:
      break;
  }
  return os;
}

std::ostream &operator<<(std::ostream &os, const Combiner &p) {
  switch (p) {
    case Combiner::Sum:
      os << "sum";
      break;
    case Combiner::Average:
      os << "average";
      break;
    case Combiner::Concat:
      os << "concat";
      break;
    default:
      break;
  }
  return os;
}

std::ostream &operator<<(std::ostream &os, const EmbeddingLayout &p) {
  switch (p) {
    case EmbeddingLayout::FeatureMajor:
      os << "FeatureMajor";
      break;
    case EmbeddingLayout::BatchMajor:
      os << "BatchMajor";
      break;
    default:
      HCTR_OWN_THROW(HugeCTR::Error_t::NotInitialized, "EmbeddingLayout is not initialized");
  }
  return os;
}

std::ostream &operator<<(std::ostream &os, const CommunicationStrategy &p) {
  switch (p) {
    case CommunicationStrategy::Uniform:
      os << "Uniform";
      break;
    case CommunicationStrategy::Hierarchical:
      os << "Hierarchical";
      break;
    default:
      HCTR_OWN_THROW(HugeCTR::Error_t::NotInitialized, "CommunicationStrategy is not initialized");
  }
  return os;
}

std::ostream &operator<<(std::ostream &os, const SortStrategy &p) {
  switch (p) {
    case SortStrategy::Radix:
      os << "Radix";
      break;
    case SortStrategy::Segmented:
      os << "Segmented";
      break;
    default:
      HCTR_OWN_THROW(HugeCTR::Error_t::NotInitialized, "SortStrategy is not initialized");
  }
  return os;
}

std::ostream &operator<<(std::ostream &os, const KeysPreprocessStrategy &p) {
  switch (p) {
    case KeysPreprocessStrategy::None:
      os << "None";
      break;
    case KeysPreprocessStrategy::AddOffset:
      os << "AddOffset";
      break;
    default:
      HCTR_OWN_THROW(HugeCTR::Error_t::NotInitialized, "KeysPreprocessStrategy is not initialized");
  }
  return os;
}

std::ostream &operator<<(std::ostream &os, const AllreduceStrategy &p) {
  switch (p) {
    case AllreduceStrategy::Dense:
      os << "Dense";
      break;
    case AllreduceStrategy::GroupDense:
      os << "GroupDense";
      break;
    default:
      HCTR_OWN_THROW(HugeCTR::Error_t::NotInitialized, "AllreduceStrategy is not initialized");
  }
  return os;
}

std::ostream &operator<<(std::ostream &os, const LookupParam &p) {
  os << "lookup_id:" << p.lookup_id << ",";
  os << "table_id:" << p.table_id << ",";
  os << "combiner:" << p.combiner << ",";
  os << "max_hotness:" << p.max_hotness << ",";
  os << "ev_size:" << p.ev_size;
  return os;
}

std::ostream &operator<<(std::ostream &os, const CompressionParam &p) {
  for (auto &[compression_strategy, table_ids] : p.compression_strategy_to_table_ids) {
    os << compression_strategy << ":";
    for (auto table_id : table_ids) {
      os << table_id << ",";
    }
  }
  return os;
}

std::vector<int> filter_dp_lookup_ids(const std::vector<LookupParam> &lookup_params,
                                      const GroupedTableParam &table_param) {
  int num_lookup = static_cast<int>(lookup_params.size());
  std::vector<int> dp_lookup_ids;
  for (int lookup_id = 0; lookup_id < num_lookup; ++lookup_id) {
    auto combiner = lookup_params[lookup_id].combiner;
    auto max_hotness = lookup_params[lookup_id].max_hotness;

    HCTR_CHECK_HINT(!((combiner == Combiner::Concat) && (max_hotness > 1)),
                    "DataParallel TablePlacementStrategy does not support Combiner::Concat && "
                    "max_hotness > 1.");

    int table_id = lookup_params[lookup_id].table_id;
    if (std::find(table_param.table_ids.begin(), table_param.table_ids.end(), table_id) ==
        table_param.table_ids.end())
      continue;
    dp_lookup_ids.push_back(lookup_id);
  }
  return dp_lookup_ids;
}

std::vector<int> filter_mp_sparse_lookup_ids(const std::vector<LookupParam> &lookup_params,
                                             const GroupedTableParam &table_param) {
  int num_lookup = static_cast<int>(lookup_params.size());
  std::vector<int> mp_sparse_lookup_ids;
  for (int lookup_id = 0; lookup_id < num_lookup; ++lookup_id) {
    int table_id = lookup_params[lookup_id].table_id;
    if (std::find(table_param.table_ids.begin(), table_param.table_ids.end(), table_id) ==
        table_param.table_ids.end())
      continue;
    auto combiner = lookup_params[lookup_id].combiner;
    if (combiner == Combiner::Sum || combiner == Combiner::Average) {
      mp_sparse_lookup_ids.push_back(lookup_id);
    }
  }
  return mp_sparse_lookup_ids;
}

std::vector<int> filter_mp_dense_lookup_ids(const std::vector<LookupParam> &lookup_params,
                                            const GroupedTableParam &table_param) {
  int num_lookup = static_cast<int>(lookup_params.size());
  std::vector<int> mp_dense_lookup_ids;
  for (int lookup_id = 0; lookup_id < num_lookup; ++lookup_id) {
    int table_id = lookup_params[lookup_id].table_id;
    if (std::find(table_param.table_ids.begin(), table_param.table_ids.end(), table_id) ==
        table_param.table_ids.end())
      continue;
    auto combiner = lookup_params[lookup_id].combiner;
    if (combiner == Combiner::Concat) {
      mp_dense_lookup_ids.push_back(lookup_id);
    }
  }
  return mp_dense_lookup_ids;
}

EmbeddingCollectionParam::EmbeddingCollectionParam(
    int num_table, int num_lookup, const std::vector<LookupParam> &lookup_params,
    const std::vector<std::vector<int>> &shard_matrix,
    const std::vector<GroupedTableParam> &grouped_table_params, int universal_batch_size,
    core23::DataType key_type, core23::DataType index_type, core23::DataType offset_type,
    core23::DataType emb_type, core23::DataType wgrad_type, EmbeddingLayout input_layout_,
    EmbeddingLayout output_layout, SortStrategy sort_strategy,
    KeysPreprocessStrategy keys_preprocess_strategy, AllreduceStrategy allreduce_strategy,
    CommunicationStrategy comm_strategy, CompressionParam compreesion_param)
    : num_table(num_table),
      num_lookup(num_lookup),
      lookup_params(lookup_params),
      shard_matrix(shard_matrix),
      grouped_table_params(grouped_table_params),
      universal_batch_size(universal_batch_size),
      key_type(key_type),
      index_type(index_type),
      offset_type(offset_type),
      emb_type(emb_type),
      wgrad_type_(wgrad_type),
      input_layout_(input_layout_),
      output_layout_(output_layout),
      sort_strategy_(sort_strategy),
      keys_preprocess_strategy_(keys_preprocess_strategy),
      allreduce_strategy_(allreduce_strategy),
      comm_strategy_(comm_strategy) {
  if (!compreesion_param.compression_strategy_to_table_ids.empty()) {
    const auto &table_ids_using_reduction =
        compreesion_param.compression_strategy_to_table_ids[CompressionStrategy::Reduction];
    const auto &table_ids_using_unique =
        compreesion_param.compression_strategy_to_table_ids[CompressionStrategy::Unique];
    std::vector<int> dup_table_ids;
    std::set_intersection(table_ids_using_reduction.begin(), table_ids_using_reduction.end(),
                          table_ids_using_unique.begin(), table_ids_using_unique.end(),
                          std::back_inserter(dup_table_ids));
    HCTR_CHECK_HINT(dup_table_ids.empty(), "Duplicate table id in different CompressionStrategy");

    std::vector<int> sum_table_ids;
    std::set_union(table_ids_using_reduction.begin(), table_ids_using_reduction.end(),
                   table_ids_using_unique.begin(), table_ids_using_unique.end(),
                   std::back_inserter(sum_table_ids));
    std::vector<int> mp_table_ids;
    for (size_t grouped_table_id = 0; grouped_table_id < grouped_table_params.size();
         ++grouped_table_id) {
      const auto &table_param = grouped_table_params[grouped_table_id];
      if (table_param.table_placement_strategy == TablePlacementStrategy::ModelParallel) {
        for (int table_id : table_param.table_ids) {
          mp_table_ids.push_back(table_id);
        }
      }
    }
    HCTR_CHECK_HINT(mp_table_ids.size() == sum_table_ids.size(),
                    "Table ids in CompressionStrategy does not match with table ids in "
                    "TablePlacementStrategy");
  }
  for (size_t grouped_table_id = 0; grouped_table_id < grouped_table_params.size();
       ++grouped_table_id) {
    const auto &table_param = grouped_table_params[grouped_table_id];
    if (table_param.table_placement_strategy == TablePlacementStrategy::DataParallel) {
      auto dp_lookup_ids = filter_dp_lookup_ids(lookup_params, table_param);
      grouped_lookup_params.emplace_back(grouped_table_id, dp_lookup_ids,
                                         EmbeddingGroupType::DataParallel);
      continue;
    }
    HCTR_CHECK_HINT(table_param.table_placement_strategy == TablePlacementStrategy::ModelParallel);

    auto mp_sparse_lookup_ids = filter_mp_sparse_lookup_ids(lookup_params, table_param);
    auto mp_dense_lookup_ids = filter_mp_dense_lookup_ids(lookup_params, table_param);

    auto compression_strategy_to_table_ids = compreesion_param.compression_strategy_to_table_ids;
    if (compression_strategy_to_table_ids.empty()) {
      // user does not specify compression param, use reduction for sparse and unique for dense in
      // default
      for (int lookup_id : mp_sparse_lookup_ids) {
        compression_strategy_to_table_ids[CompressionStrategy::Reduction].insert(
            lookup_params[lookup_id].table_id);
      }
      for (int lookup_id : mp_dense_lookup_ids) {
        compression_strategy_to_table_ids[CompressionStrategy::Unique].insert(
            lookup_params[lookup_id].table_id);
      }
    }
    const auto &table_ids_using_reduction =
        compression_strategy_to_table_ids[CompressionStrategy::Reduction];
    const auto &table_ids_using_unique =
        compression_strategy_to_table_ids[CompressionStrategy::Unique];

    // do sanity check
    for (int lookup_id : mp_dense_lookup_ids) {
      HCTR_CHECK_HINT(table_ids_using_unique.find(lookup_params[lookup_id].table_id) !=
                          table_ids_using_unique.end(),
                      "Unique Compression can only work for Dense");
    }
    for (int lookup_id : mp_dense_lookup_ids) {
      int table_id = lookup_params[lookup_id].table_id;
      HCTR_CHECK_HINT(
          table_ids_using_unique.find(table_id) != table_ids_using_unique.end() ||
              table_ids_using_reduction.find(table_id) != table_ids_using_reduction.end(),
          "Unique or Reduction Compression can work for Sparse");
    }

    auto filter_and_add_grouped_lookup_by_compression_strategy =
        [&](const std::vector<int> &lookup_ids, const CompressionStrategy &compression_strategy) {
          std::vector<int> filtered_lookup_ids;
          std::copy_if(lookup_ids.begin(), lookup_ids.end(),
                       std::back_inserter(filtered_lookup_ids), [&](int lookup_id) {
                         auto table_id = lookup_params[lookup_id].table_id;
                         auto &table_ids = compression_strategy_to_table_ids[compression_strategy];
                         return table_ids.find(table_id) != table_ids.end();
                       });
          return filtered_lookup_ids;
        };
    if (auto mp_sparse_lookup_ids_using_reduction =
            filter_and_add_grouped_lookup_by_compression_strategy(mp_sparse_lookup_ids,
                                                                  CompressionStrategy::Reduction);
        !mp_sparse_lookup_ids_using_reduction.empty()) {
      grouped_lookup_params.emplace_back(grouped_table_id, mp_sparse_lookup_ids_using_reduction,
                                         EmbeddingGroupType::SparseModelParallel);
    }
    if (auto mp_sparse_lookup_ids_using_unique =
            filter_and_add_grouped_lookup_by_compression_strategy(mp_sparse_lookup_ids,
                                                                  CompressionStrategy::Unique);
        !mp_sparse_lookup_ids_using_unique.empty()) {
      std::unordered_map<int, std::vector<int>> ev_size_to_lookup_id_dict;
      for (auto lookup_id : mp_sparse_lookup_ids_using_unique) {
        ev_size_to_lookup_id_dict[lookup_params[lookup_id].ev_size].push_back(lookup_id);
      }
      for (auto &[_, lookup_ids_with_same_ev_size] : ev_size_to_lookup_id_dict) {
        grouped_lookup_params.emplace_back(grouped_table_id, lookup_ids_with_same_ev_size,
                                           EmbeddingGroupType::DenseModelParallelWithReduction);
      }
    }
    if (!mp_dense_lookup_ids.empty()) {
      auto mp_dense_lookup_ids_using_unique = filter_and_add_grouped_lookup_by_compression_strategy(
          mp_dense_lookup_ids, CompressionStrategy::Unique);
      HCTR_CHECK_HINT(mp_dense_lookup_ids == mp_dense_lookup_ids_using_unique,
                      "CompressionStrategy is wrong for dense embedding");
      std::unordered_map<int, std::vector<int>> ev_size_to_lookup_id_dict;
      for (auto lookup_id : mp_dense_lookup_ids) {
        ev_size_to_lookup_id_dict[lookup_params[lookup_id].ev_size].push_back(lookup_id);
      }
      for (auto &[_, lookup_ids_with_same_ev_size] : ev_size_to_lookup_id_dict) {
        grouped_lookup_params.emplace_back(grouped_table_id, lookup_ids_with_same_ev_size,
                                           EmbeddingGroupType::DenseModelParallel);
      }
    }
  }
}

void EmbeddingOutputAttr::init(std::shared_ptr<CoreResourceManager> core,
                               const EmbeddingCollectionParam &ebc_param) {
  this->num_lookup = ebc_param.num_lookup;
  const auto &lookup_params = ebc_param.lookup_params;
  h_id_to_ev_size.clear();
  h_id_to_combiner.clear();
  h_id_to_ev_start_indices = {0};

  for (int lookup_id = 0; lookup_id < ebc_param.num_lookup; ++lookup_id) {
    int ev_size = lookup_params[lookup_id].ev_size;
    h_id_to_ev_size.push_back(ev_size);
    char combiner = static_cast<char>(lookup_params[lookup_id].combiner);
    h_id_to_combiner.push_back(combiner);

    h_id_to_ev_start_indices.push_back((combiner == static_cast<char>(Combiner::Concat))
                                           ? ebc_param.lookup_params[lookup_id].max_hotness *
                                                 ev_size
                                           : ev_size);
  }

  std::partial_sum(h_id_to_ev_start_indices.begin(), h_id_to_ev_start_indices.end(),
                   h_id_to_ev_start_indices.begin());

  HugeCTR::CudaDeviceContext context(core->get_device_id());
  core23::Device device(core23::DeviceType::GPU, core->get_device_id());
  core23::BufferParams buffer_params;
  buffer_params.unitary = false;
  core23::TensorParams params = core23::TensorParams().device(device).buffer_params(buffer_params);

  this->id_to_ev_size = core23::Tensor(params.shape({static_cast<int64_t>(h_id_to_ev_size.size())})
                                           .data_type(core23::ScalarType::Int32));
  core23::copy_sync(this->id_to_ev_size, h_id_to_ev_size);

  this->id_to_ev_start_indices =
      core23::Tensor(params.shape({static_cast<int64_t>(h_id_to_ev_start_indices.size())})
                         .data_type(core23::ScalarType::Int32));
  core23::copy_sync(this->id_to_ev_start_indices, h_id_to_ev_start_indices);

  this->id_to_combiner =
      core23::Tensor(params.shape({static_cast<int64_t>(h_id_to_combiner.size())})
                         .data_type(core23::ScalarType::Char));
  core23::copy_sync(this->id_to_combiner, h_id_to_combiner);

  this->num_elements_per_sample = h_id_to_ev_start_indices.back();

  this->layout = ebc_param.output_layout_;
  this->max_ev_size = !h_id_to_ev_size.empty()
                          ? *std::max_element(h_id_to_ev_size.begin(), h_id_to_ev_size.end())
                          : 0;
  this->is_ragged =
      (h_id_to_ev_size.size() == 0)
          ? false
          : std::equal(h_id_to_ev_size.begin() + 1, h_id_to_ev_size.end(), h_id_to_ev_size.begin());

  bool is_aligned = true;
  for (auto ev_size : h_id_to_ev_size) {
    if (ev_size % 4 != 0) is_aligned = false;
  }
  this->is_aligned = is_aligned;

  this->type = ebc_param.emb_type;
}

void EmbeddingOutputAttr::update_mutable_data(std::shared_ptr<CoreResourceManager> core,
                                              const EmbeddingCollectionParam &ebc_param) const {
  h_id_to_hotness.clear();

  HugeCTR::CudaDeviceContext context(core->get_device_id());
  const auto &lookup_params = ebc_param.lookup_params;

  size_t num_gpus = core->get_global_gpu_count();

  HCTR_CHECK_HINT(ebc_param.shard_matrix.size() == num_gpus,
                  "shard matrix should contain num_gpus row.");

  for (int lookup_id = 0; lookup_id < this->num_lookup; ++lookup_id) {
    int max_hotness = lookup_params[lookup_id].max_hotness;

    h_id_to_hotness.push_back(max_hotness);
  }
  hotness_sum = std::accumulate(h_id_to_hotness.begin(), h_id_to_hotness.end(), 0);
}

std::vector<int> get_lookup_id_table_id(const EmbeddingCollectionParam &ebc_param,
                                        size_t grouped_id, int gpu_id) {
  const auto &lookup_params = ebc_param.lookup_params;

  std::vector<int> lookup_id_to_table_id;
  for (int lookup_id = 0; lookup_id < ebc_param.num_lookup; ++lookup_id) {
    if (!ebc_param.has_table_shard(gpu_id, grouped_id, lookup_id)) continue;

    int table_id = lookup_params[lookup_id].table_id;
    lookup_id_to_table_id.push_back(table_id);
  }
  return lookup_id_to_table_id;
}

std::vector<int> sort_lookup_ids(const std::vector<int> &lookup_id_to_table_id) {
  int num_lookup = lookup_id_to_table_id.size();

  std::vector<int> sorted_lookup_ids(num_lookup);
  std::iota(sorted_lookup_ids.begin(), sorted_lookup_ids.end(), 0);
  std::sort(sorted_lookup_ids.begin(), sorted_lookup_ids.end(),
            [&](int l, int r) { return lookup_id_to_table_id[l] < lookup_id_to_table_id[r]; });

  return sorted_lookup_ids;
}

std::vector<int> get_sorted_table_ids(const std::vector<int> &sorted_lookup_ids,
                                      const std::vector<int> &lookup_id_to_table_id) {
  std::vector<int> sorted_table_ids;
  std::transform(sorted_lookup_ids.begin(), sorted_lookup_ids.end(),
                 std::back_inserter(sorted_table_ids),
                 [&](int idx) { return lookup_id_to_table_id[idx]; });

  return sorted_table_ids;
}

std::vector<int> deduplicate_sorted_table_ids(const std::vector<int> &sorted_table_ids) {
  std::vector<int> unique_table_ids{sorted_table_ids.begin(), sorted_table_ids.end()};
  auto last = std::unique(unique_table_ids.begin(), unique_table_ids.end());
  unique_table_ids.erase(last, unique_table_ids.end());
  return unique_table_ids;
}

std::vector<int> get_table_id_to_ev_size(std::shared_ptr<CoreResourceManager> core,
                                         const EmbeddingCollectionParam &ebc_param) {
  const auto &lookup_params = ebc_param.lookup_params;

  std::vector<int> h_table_id_to_ev_size;
  h_table_id_to_ev_size.resize(ebc_param.num_table);
  for (int lookup_id = 0; lookup_id < ebc_param.num_lookup; ++lookup_id) {
    int table_id = lookup_params[lookup_id].table_id;
    h_table_id_to_ev_size[table_id] = lookup_params[lookup_id].ev_size;
  }
  return h_table_id_to_ev_size;
}

void WgradAttr::init(std::shared_ptr<CoreResourceManager> core,
                     const EmbeddingCollectionParam &ebc_param, size_t grouped_id) {
  int gpu_id = core->get_global_gpu_id();
  auto h_lookup_id_to_table_id = get_lookup_id_table_id(ebc_param, grouped_id, gpu_id);
  auto h_sorted_lookup_ids = sort_lookup_ids(h_lookup_id_to_table_id);
  auto h_sorted_table_ids = get_sorted_table_ids(h_sorted_lookup_ids, h_lookup_id_to_table_id);
  this->h_sorted_unique_table_ids = deduplicate_sorted_table_ids(h_sorted_table_ids);
  auto h_table_id_to_ev_size = get_table_id_to_ev_size(core, ebc_param);

  this->num_table = static_cast<int>(this->h_sorted_unique_table_ids.size());
  this->num_lookup = h_lookup_id_to_table_id.size();
  this->type = ebc_param.wgrad_type_;

  HugeCTR::CudaDeviceContext context(core->get_device_id());
  core23::Device device(core23::DeviceType::GPU, core->get_device_id());
  core23::BufferParams buffer_params;
  buffer_params.unitary = false;
  core23::TensorParams params = core23::TensorParams().device(device).buffer_params(buffer_params);

  const auto &lookup_params = ebc_param.lookup_params;
  std::vector<int> h_ev_size_list;
  h_ev_size_list.clear();

  for (int lookup_id = 0; lookup_id < ebc_param.num_lookup; ++lookup_id) {
    h_ev_size_list.push_back(lookup_params[lookup_id].ev_size);
  }

  // FIX:global ev size same or local ev size same? for now is global ev size same
  // I think is include dp and mp.
  if (h_ev_size_list.size() > 0) {
    same_ev_size = h_ev_size_list[0];
    is_same_ev_size = true;
    for (size_t i = 1; i < h_ev_size_list.size(); ++i) {
      if (h_ev_size_list[i] != same_ev_size) {
        is_same_ev_size = false;
        same_ev_size = 0;
        break;
      }
    }
  }

  this->lookup_id_to_table_ids =
      core23::Tensor(params.shape({static_cast<int64_t>(h_lookup_id_to_table_id.size())})
                         .data_type(core23::ScalarType::Int32));
  core23::copy_sync(this->lookup_id_to_table_ids, h_lookup_id_to_table_id);

  this->sorted_lookup_ids =
      core23::Tensor(params.shape({static_cast<int64_t>(h_sorted_lookup_ids.size())})
                         .data_type(core23::ScalarType::Int32));
  core23::copy_sync(this->sorted_lookup_ids, h_sorted_lookup_ids);

  this->sorted_table_ids =
      core23::Tensor(params.shape({static_cast<int64_t>(h_sorted_table_ids.size())})
                         .data_type(core23::ScalarType::Int32));
  core23::copy_sync(this->sorted_table_ids, h_sorted_table_ids);

  this->sorted_unique_table_ids =
      core23::Tensor(params.shape({static_cast<int64_t>(h_sorted_unique_table_ids.size())})
                         .data_type(core23::ScalarType::Int32));
  core23::copy_sync(this->sorted_unique_table_ids, h_sorted_unique_table_ids);

  this->table_id_to_ev_size =
      core23::Tensor(params.shape({static_cast<int64_t>(h_table_id_to_ev_size.size())})
                         .data_type(core23::ScalarType::Int32));
  core23::copy_sync(this->table_id_to_ev_size, h_table_id_to_ev_size);
}

std::vector<int> get_wgrad_max_num_keys(const EmbeddingCollectionParam &ebc_param,
                                        size_t grouped_id, int gpu_id) {
  const auto &lookup_params = ebc_param.lookup_params;

  std::vector<int> local_max_hotness_list;
  for (int lookup_id = 0; lookup_id < ebc_param.num_lookup; ++lookup_id) {
    if (!ebc_param.has_table_shard(gpu_id, grouped_id, lookup_id)) continue;

    int max_hotness = lookup_params[lookup_id].max_hotness;
    local_max_hotness_list.push_back(max_hotness);
  }
  return local_max_hotness_list;
}

std::vector<int> get_wgrad_ev_size(const EmbeddingCollectionParam &ebc_param, size_t grouped_id,
                                   int gpu_id) {
  const auto &lookup_params = ebc_param.lookup_params;
  std::vector<int> local_ev_size_list;
  for (int lookup_id = 0; lookup_id < ebc_param.num_lookup; ++lookup_id) {
    if (!ebc_param.has_table_shard(gpu_id, grouped_id, lookup_id)) continue;

    int ev_size = lookup_params[lookup_id].ev_size;
    local_ev_size_list.push_back(ev_size);
  }
  return local_ev_size_list;
}

std::vector<int> get_allreduce_buffer_num_keys(
    const std::vector<int> &unique_table_ids, const std::vector<int> &table_id_to_vocabulary_size) {
  std::vector<int> vocabulary_size_list;
  for (auto table_id : unique_table_ids) {
    vocabulary_size_list.push_back(table_id_to_vocabulary_size[table_id]);
  }
  return vocabulary_size_list;
}

void Wgrad::init_attr(const EmbeddingCollectionParam &ebc_param, WgradAttr wgrad_attr,
                      std::shared_ptr<CoreResourceManager> &core) {
  HugeCTR::CudaDeviceContext context(core->get_device_id());
  core23::Device device(core23::DeviceType::GPU, core->get_device_id());
  core23::TensorParams params = core23::TensorParams().device(device);
  this->attr = wgrad_attr;
  this->attr.type = ebc_param.wgrad_type_;
  this->attr.key_type = ebc_param.key_type;
  this->num_unique_keys = core23::Tensor(params.shape({1}).data_type(core23::ScalarType::UInt64));
  this->attr_allocated = true;
}

void Wgrad::init_data(std::shared_ptr<CoreResourceManager> core, uint32_t h_unique_key_num,
                      int max_ev_size) {
  HCTR_CHECK_HINT(this->attr_allocated, "Wgrad's attr should be inited before init data");
  HugeCTR::CudaDeviceContext context(core->get_device_id());
  core23::Device device(core23::DeviceType::GPU, core->get_device_id());
  core23::TensorParams params = core23::TensorParams().device(device);
  this->unique_keys = core23::Tensor(
      params.shape({static_cast<int64_t>(h_unique_key_num)}).data_type(this->attr.key_type));
  this->table_ids = core23::Tensor(
      params.shape({static_cast<int64_t>(h_unique_key_num)}).data_type(core23::ScalarType::Int32));
  this->data = core23::Tensor(params.shape({static_cast<int64_t>(h_unique_key_num) * max_ev_size})
                                  .data_type(this->attr.type));
  this->ev_start_indices = core23::Tensor(params.shape({static_cast<int64_t>(h_unique_key_num) + 1})
                                              .data_type(core23::ScalarType::UInt32));

  this->h_unique_keys = h_unique_key_num;
  this->is_dynamic_allocate = true;
  this->data_allocated = true;
  this->indices_allocated = true;
  return;
}

void DynamicWgradInitiazlier::init_indices_and_data(Wgrad &wgrad, size_t h_num_key,
                                                    core23::Tensor key_flag_buffer, int max_ev_size,
                                                    std::shared_ptr<CoreResourceManager> core,
                                                    cudaStream_t stream) {
  HCTR_LIB_THROW(cudaStreamSynchronize(stream));
  uint32_t h_unique_key_num = 0;
  HCTR_LIB_THROW(cudaMemcpy(&h_unique_key_num, key_flag_buffer.data<uint32_t>() + h_num_key - 1,
                            sizeof(uint32_t), cudaMemcpyDeviceToHost));
  wgrad.init_data(core, h_unique_key_num, max_ev_size);
  HCTR_LIB_THROW(cudaStreamSynchronize(stream));
  return;
}

WgradInitializer &WgradInitializer::init(Wgrad &other) {
  this->wgrad = &other;
  this->wgrad->init_attr(ebc_param, wgrad_attr, core);
  return *this;
}
WgradInitializer &WgradInitializer::init_indices() {
  HCTR_CHECK_HINT(wgrad->attr_allocated, "Wgrad's attr should be inited before init indices");
  int gpu_id = core->get_global_gpu_id();

  auto local_max_hotness_list = get_wgrad_max_num_keys(ebc_param, grouped_id, gpu_id);
  auto local_ev_size_list = get_wgrad_ev_size(ebc_param, grouped_id, gpu_id);

  int max_num_keys =
      std::accumulate(local_max_hotness_list.begin(), local_max_hotness_list.end(), 0);

  auto key_type = ebc_param.key_type;
  int batch_size = ebc_param.universal_batch_size;

  HugeCTR::CudaDeviceContext context(core->get_device_id());

  core23::Device device(core23::DeviceType::GPU, core->get_device_id());
  core23::TensorParams params = core23::TensorParams().device(device);
  wgrad->attr.type = ebc_param.wgrad_type_;
  wgrad->unique_keys =
      core23::Tensor(params.shape({batch_size * max_num_keys}).data_type(key_type));
  wgrad->table_ids = core23::Tensor(
      params.shape({batch_size * max_num_keys}).data_type(core23::ScalarType::Int32));
  wgrad->ev_start_indices = core23::Tensor(
      params.shape({batch_size * max_num_keys + 1}).data_type(core23::ScalarType::UInt32));
  wgrad->indices_allocated = true;
  return *this;
}

WgradInitializer &WgradInitializer::init_data() {
  HCTR_CHECK_HINT(wgrad->attr_allocated, "Wgrad's attr should be inited before init data");
  int gpu_id = core->get_global_gpu_id();

  auto local_max_hotness_list = get_wgrad_max_num_keys(ebc_param, grouped_id, gpu_id);
  auto local_ev_size_list = get_wgrad_ev_size(ebc_param, grouped_id, gpu_id);

  int batch_size = ebc_param.universal_batch_size;

  HugeCTR::CudaDeviceContext context(core->get_device_id());

  int64_t max_buffer_size = 0;
  for (size_t i = 0; i < local_max_hotness_list.size(); ++i) {
    max_buffer_size += local_max_hotness_list[i] * local_ev_size_list[i];
  }
  max_buffer_size *= batch_size;
  core23::Device device(core23::DeviceType::GPU, core->get_device_id());
  core23::TensorParams params = core23::TensorParams().device(device);

  double wgrad_unique_ratio = get_wgrad_unique_ratio();
  int64_t num_elements =
      static_cast<int64_t>(wgrad_unique_ratio * static_cast<double>(max_buffer_size));

  wgrad->data = core23::Tensor(params.shape({num_elements}).data_type(wgrad->attr.type));
  wgrad->data_allocated = true;
  return *this;
}

AllreduceWgradInitializer &AllreduceWgradInitializer::init(Wgrad &other) {
  this->wgrad = &other;
  this->wgrad->init_attr(ebc_param, wgrad_attr, core);
  return *this;
}

AllreduceWgradInitializer &AllreduceWgradInitializer::init_indices() {
  HCTR_CHECK_HINT(wgrad->attr_allocated, "Wgrad's attr should be inited before init indices");
  int gpu_id = core->get_global_gpu_id();

  std::vector<int> h_local_ev_size_list = get_wgrad_ev_size(ebc_param, grouped_id, gpu_id);
  std::vector<int> h_unique_table_ids(wgrad_attr.sorted_unique_table_ids.num_elements());
  core23::copy_sync(h_unique_table_ids, wgrad_attr.sorted_unique_table_ids);

  std::vector<int> h_local_num_keys_list =
      get_allreduce_buffer_num_keys(h_unique_table_ids, table_id_to_vocabulary_size);
  auto key_type = ebc_param.key_type;

  int max_num_keys = std::accumulate(h_local_num_keys_list.begin(), h_local_num_keys_list.end(), 0);

  HugeCTR::CudaDeviceContext context(core->get_device_id());

  core23::Device device(core23::DeviceType::GPU, core->get_device_id());
  core23::TensorParams params = core23::TensorParams().device(device);

  wgrad->unique_keys = core23::Tensor(params.shape({max_num_keys}).data_type(key_type));
  wgrad->table_ids =
      core23::Tensor(params.shape({max_num_keys}).data_type(core23::ScalarType::Int32));
  wgrad->ev_start_indices =
      core23::Tensor(params.shape({max_num_keys + 1}).data_type(core23::ScalarType::UInt32));

  // TODO: move those initialization on GPU
  DISPATCH_INTEGRAL_FUNCTION_CORE23(key_type.type(), key_t, [&] {
    std::vector<key_t> h_unique_keys;
    size_t num_keys = 0;
    for (size_t i = 0; i < h_local_num_keys_list.size(); ++i) {
      num_keys += h_local_num_keys_list[i];
    }
    h_unique_keys.resize(num_keys);
    std::iota(h_unique_keys.begin(), h_unique_keys.end(), 0);
    core23::copy_sync(wgrad->unique_keys, h_unique_keys);
    std::vector<size_t> h_num_unqiue_keys{h_unique_keys.size()};
    core23::copy_sync(wgrad->num_unique_keys.data(), h_num_unqiue_keys.data(),
                      wgrad->num_unique_keys.num_bytes(), wgrad->num_unique_keys.device(),
                      core23::DeviceType::CPU);
  });

  std::vector<int> h_table_ids;
  for (size_t i = 0; i < h_local_num_keys_list.size(); ++i) {
    h_table_ids.insert(h_table_ids.end(), h_local_num_keys_list[i], h_unique_table_ids[i]);
  }
  core23::copy_sync(wgrad->table_ids.data(), h_table_ids.data(), wgrad->table_ids.num_bytes(),
                    wgrad->table_ids.device(), core23::DeviceType::CPU);
  {
    std::vector<uint32_t> h_ev_start_indices;
    uint32_t cnt = 0;
    for (size_t i = 0; i < h_local_num_keys_list.size(); ++i) {
      int ev_size = h_local_ev_size_list[i];
      int num_keys = h_local_num_keys_list[i];
      for (int ik = 0; ik < num_keys; ++ik) {
        h_ev_start_indices.push_back(cnt);
        cnt += ev_size;
      }
    }
    h_ev_start_indices.push_back(cnt);
    core23::copy_sync(wgrad->ev_start_indices.data(), h_ev_start_indices.data(),
                      wgrad->ev_start_indices.num_bytes(), wgrad->ev_start_indices.device(),
                      core23::DeviceType::CPU);
  }
  wgrad->indices_allocated = true;
  return *this;
}

AllreduceWgradInitializer &AllreduceWgradInitializer::init_data(
    bool grouped, const core23::BufferChannel &buffer_channel) {
  HCTR_CHECK_HINT(wgrad->attr_allocated, "Wgrad's attr should be inited before init data");
  wgrad->attr = wgrad_attr;
  int gpu_id = core->get_global_gpu_id();
  std::vector<int> h_local_ev_size_list = get_wgrad_ev_size(ebc_param, grouped_id, gpu_id);
  std::vector<int> h_unique_table_ids(wgrad_attr.sorted_unique_table_ids.num_elements());
  core23::copy_sync(h_unique_table_ids, wgrad_attr.sorted_unique_table_ids);
  std::vector<int> h_local_num_keys_list =
      get_allreduce_buffer_num_keys(h_unique_table_ids, table_id_to_vocabulary_size);

  HugeCTR::CudaDeviceContext context(core->get_device_id());

  int64_t max_buffer_size = 0;
  for (size_t i = 0; i < h_local_num_keys_list.size(); ++i) {
    max_buffer_size += h_local_num_keys_list[i] * h_local_ev_size_list[i];
  }

  core23::Device device(core23::DeviceType::GPU, core->get_device_id());
  core23::TensorParams wgrads_params = core23::TensorParams().device(device);
  int alignment_num = 32;
  if (grouped) {
    // out-of-place modifications
    wgrads_params = wgrads_params.alignment(alignment_num).buffer_channel(buffer_channel);
  }
  wgrad->data = core23::Tensor(wgrads_params.shape({max_buffer_size}).data_type(wgrad->attr.type));
  wgrad->data_allocated = true;

  return *this;
}

double get_dense_unique_ratio() {
  double dense_unique_ratio = 1.0f;

  const char *const dense_unique_ratio_env = std::getenv("DENSE_UNIQUE_RATIO");
  if (dense_unique_ratio_env != nullptr) {
    dense_unique_ratio = std::atof(dense_unique_ratio_env);
  }
  return dense_unique_ratio;
}

double get_wgrad_unique_ratio() {
  double wgrad_unique_ratio = 1.0f;

  const char *const dense_unique_ratio_env = std::getenv("WGRAD_UNIQUE_RATIO");
  if (dense_unique_ratio_env != nullptr) {
    wgrad_unique_ratio = std::atof(dense_unique_ratio_env);
  }
  return wgrad_unique_ratio;
}
}  // namespace embedding
