////////////////////////////////////////////////////////////////////////////////
// Copyright (c) 2018, Lawrence Livermore National Security, LLC.  Produced at the
// Lawrence Livermore National Laboratory in collaboration with University of
// Illinois Urbana-Champaign.
//
// Written by the LBANN Research Team (N. Dryden, N. Maruyama, et al.) listed in
// the CONTRIBUTORS file. <lbann-dev@llnl.gov>
//
// LLNL-CODE-756777.
// All rights reserved.
//
// This file is part of Aluminum GPU-aware Communication Library. For details, see
// http://software.llnl.gov/Aluminum or https://github.com/LLNL/Aluminum.
//
// Licensed under the Apache License, Version 2.0 (the "Licensee"); 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.
////////////////////////////////////////////////////////////////////////////////

#pragma once

#include <Al_config.hpp>

#include <array>
#include <vector>
#include <mutex>
#include <unordered_map>
#include "aluminum/base.hpp"
#include "aluminum/utils/meta.hpp"

#ifdef AL_HAS_CUDA
#include "aluminum/cuda/cuda.hpp"
#endif

namespace Al {
namespace internal {

namespace details {

/** Return the number of bins the caching allocator will use. */
template <typename Params>
constexpr size_t num_bins() {
  size_t num = num_pow2s_between<size_t, Params::min_bin_size, Params::max_bin_size>();
  // Count all the remaining bins that are not powers of 2. */
  for (float bin_size = Params::min_bin_size;
       bin_size <= Params::max_bin_size;
       bin_size *= Params::bin_growth) {
    if (!is_pow2(static_cast<size_t>(bin_size))) {
      ++num;
    }
  }
  return num;
}

/**
 * Return the bin sizes for the caching allocator in sorted order.
 *
 * The bins consist of all values generated by the parameters, plus
 * all powers of 2 in the range, if not present.
 */
template <typename Params>
constexpr std::array<size_t, num_bins<Params>()> get_bin_sizes() {
  // Note: To be used in a constexpr, this function requires C++17.
  // (Because operator[] for std::array is not constexpr until then.)
  std::array<size_t, num_bins<Params>()> bin_sizes{};
  auto pow2s = pow2_ar<size_t, Params::min_bin_size, Params::max_bin_size>();
  size_t pow2_idx = 0;  // Last power of 2 we used.
  size_t bin_idx = 0;  // Index into bin_sizes.
  for (float bin_size = Params::min_bin_size;
       bin_size <= Params::max_bin_size;
       bin_size *= Params::bin_growth) {
    size_t bin_size_st = static_cast<size_t>(bin_size);
    // Add all powers of 2 <= the bin size.
    for (; pow2_idx < pow2s.size() && pow2s[pow2_idx] <= bin_size_st; ++pow2_idx) {
      bin_sizes[bin_idx++] = pow2s[pow2_idx];
    }
    if (!is_pow2(bin_size_st)) {
      bin_sizes[bin_idx++] = bin_size;
    }
  }
  // Add remaining powers of 2.
  for (; pow2_idx < pow2s.size(); ++pow2_idx) {
    bin_sizes[bin_idx++] = pow2s[pow2_idx];
  }
  return bin_sizes;
}

}  // namespace details

/** Different kinds of memory. */
enum class MemoryType {
  /** Memory allocated on the host. */
  HOST,
#ifdef AL_HAS_CUDA
  /** Memory allocated on the GPU. */
  CUDA,
  /** Memory allocated on the host, pinned and registered with CUDA. */
  CUDA_PINNED_HOST
#endif
};

/** Default parameters for CachingAllocator. */
struct CachingAllocatorDefaultParams {
  /** Growth factor for bins. */
  static constexpr float bin_growth = 1.6f;
  /** Minimum bin size. */
  static constexpr size_t min_bin_size = 1;
  /** Maximum bin size. */
  static constexpr size_t max_bin_size = 1<<26;
};

/**
 * Implements a bucketed caching allocator that synchronizes with compute
 * streams.
 *
 * This is thread-safe.
 *
 * @tparam MemType Type of memory the pool manages.
 * @tparam MemoryAllocator Underlying allocator for memory.
 * @tparam Params Allocator parameters (see CachingAllocatorDefaultParams).
 */
template <MemoryType MemType, typename MemoryAllocator,
          typename Params = CachingAllocatorDefaultParams>
class CachingAllocator {
  // Implementation details:
  // MemoryAllocator must provide allocate(size in bytes) and
  // deallocate(ptr) methods.
  // Default implementation assumes CPU streams (i.e., does not do
  // synchronization).
  // Specialize the class template if this is not sufficient.
public:
  CachingAllocator() {
    for (size_t i = 0; i < bin_sizes.size(); ++i) {
      free_data.emplace_back();
    }
  }

  ~CachingAllocator() {
    clear();
  }

  /** Allocate space for size instances of type T. */
  template <typename T>
  T* allocate(size_t size) {
    const size_t size_bytes = sizeof(T) * size;
    size_t bin = get_bin(size_bytes);
    void* mem = nullptr;
    std::lock_guard<std::mutex> lg(mutex);
    // Large allocations are not cached.
    if (bin == INVALID_BIN) {
      mem = allocator.allocate(size_bytes);
    } else {
      // Check if we have cached memory.
      if (free_data[bin].empty()) {
        mem = allocator.allocate(bin_sizes[bin]);
      } else {
        mem = free_data[bin].back();
        free_data[bin].pop_back();
      }
    }
    alloc_to_bin[mem] = bin;
    return reinterpret_cast<T*>(mem);
  }

  /** Return ptr to the cache for reuse. */
  template <typename T>
  void release(T* ptr) {
    std::lock_guard<std::mutex> lg(mutex);
    auto i = alloc_to_bin.find(ptr);
    if (i == alloc_to_bin.end()) {
      throw_al_exception("Tried to free unknown ptr");
    }
    size_t bin = i->second;
    alloc_to_bin.erase(i);
    if (bin == INVALID_BIN) {
      allocator.deallocate(ptr);
    } else {
      free_data[bin].push_back(ptr);
    }
  }

  /** Release all cached memory. */
  void clear() {
    std::lock_guard<std::mutex> lg(mutex);
    for (size_t bin = 0; bin < bin_sizes.size(); ++bin) {
      for (auto&& ptr : free_data[bin]) {
        allocator.deallocate(ptr);
      }
      free_data[bin].clear();
    }
  }

private:
  /** Index for an invalid bin. */
  static constexpr size_t INVALID_BIN = static_cast<size_t>(-1);

  /** Size, in bytes, of each cache bin, in sorted order. */
  static constexpr std::array<size_t, details::num_bins<Params>()> bin_sizes =
    details::get_bin_sizes<Params>();

  /**
   * Data available to allocate.
   *
   * Each entry is a bin, and each bin contains free data for that bin.
   */
  std::vector<std::vector<void*>> free_data;

  /** Map allocated pointers to their corresponding bin index. */
  std::unordered_map<void*, size_t> alloc_to_bin;

  /** Underlying allocator. */
  MemoryAllocator allocator;

  /** Protect all accesses. */
  std::mutex mutex;

  /** Return the index of the smallest bin size >= size. */
  inline size_t get_bin(const size_t size) {
    // We use branchless binary search here and hope the compiler unrolls.
    // See: https://arxiv.org/pdf/1509.05053.pdf
    // For some cases, a counting-based linear search may be faster.
    // See: https://dirtyhandscoding.github.io/posts/performance-comparison-linear-search-vs-binary-search.html
    // Note that since we know the bins at compile time, we may be able
    // to do better.
    size_t base = 0;
    size_t n = bin_sizes.size();
    while (n > 1) {
      const size_t half = n / 2;
      base = (bin_sizes[base + half] < size) ? base + half : base;
      n -= half;
    }
    size_t r = (bin_sizes[base] < size) + base;
    return (r >= bin_sizes.size()) ? INVALID_BIN : r;
  }
};

}  // namespace internal
}  // namespace Al
