core/perf_test/test_mempool.cpp

/*
//@HEADER
// ************************************************************************
//
//                        Kokkos v. 3.0
//       Copyright (2020) National Technology & Engineering
//               Solutions of Sandia, LLC (NTESS).
//
// Under the terms of Contract DE-NA0003525 with NTESS,
// the U.S. Government retains certain rights in this software.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
// 1. Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright
// notice, this list of conditions and the following disclaimer in the
// documentation and/or other materials provided with the distribution.
//
// 3. Neither the name of the Corporation nor the names of the
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY NTESS "AS IS" AND ANY
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NTESS OR THE
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
// Questions? Contact Christian R. Trott (crtrott@sandia.gov)
//
// ************************************************************************
//@HEADER
*/

#include <cstdio>
#include <cstring>
#include <cstdlib>
#include <limits>

#include <Kokkos_Core.hpp>
#include <impl/Kokkos_Timer.hpp>

using ExecSpace   = Kokkos::DefaultExecutionSpace;
using MemorySpace = Kokkos::DefaultExecutionSpace::memory_space;

using MemoryPool = Kokkos::MemoryPool<ExecSpace>;

struct TestFunctor {
  using ptrs_type = Kokkos::View<uintptr_t*, ExecSpace>;

  enum : unsigned { chunk = 32 };

  MemoryPool pool;
  ptrs_type ptrs;
  unsigned chunk_span;
  unsigned fill_stride;
  unsigned range_iter;
  unsigned repeat_inner;

  TestFunctor(size_t total_alloc_size, unsigned min_superblock_size,
              unsigned number_alloc, unsigned arg_stride_alloc,
              unsigned arg_chunk_span, unsigned arg_repeat)
      : pool(), ptrs(), chunk_span(0), fill_stride(0), repeat_inner(0) {
    MemorySpace m;

    const unsigned min_block_size = chunk;
    const unsigned max_block_size = chunk * arg_chunk_span;
    pool = MemoryPool(m, total_alloc_size, min_block_size, max_block_size,
                      min_superblock_size);

    ptrs         = ptrs_type(Kokkos::view_alloc(m, "ptrs"), number_alloc);
    fill_stride  = arg_stride_alloc;
    chunk_span   = arg_chunk_span;
    range_iter   = fill_stride * number_alloc;
    repeat_inner = arg_repeat;
  }

  //----------------------------------------

  using value_type = long;

  //----------------------------------------

  struct TagFill {};

  KOKKOS_INLINE_FUNCTION
  void operator()(TagFill, int i, value_type& update) const noexcept {
    if (0 == i % fill_stride) {
      const int j = i / fill_stride;

      const unsigned size_alloc = chunk * (1 + (j % chunk_span));

      ptrs(j) = (uintptr_t)pool.allocate(size_alloc);

      if (ptrs(j)) ++update;
    }
  }

  bool test_fill() {
    using policy = Kokkos::RangePolicy<ExecSpace, TagFill>;

    long result = 0;

    Kokkos::parallel_reduce(policy(0, range_iter), *this, result);

    if (result == long(ptrs.extent(0))) return true;
    pool.print_state(std::cerr);
    return false;
  }

  //----------------------------------------

  struct TagDel {};

  KOKKOS_INLINE_FUNCTION
  void operator()(TagDel, int i) const noexcept {
    if (0 == i % fill_stride) {
      const int j = i / fill_stride;

      const unsigned size_alloc = chunk * (1 + (j % chunk_span));

      pool.deallocate((void*)ptrs(j), size_alloc);
    }
  }

  void test_del() {
    using policy = Kokkos::RangePolicy<ExecSpace, TagDel>;

    Kokkos::parallel_for(policy(0, range_iter), *this);
    Kokkos::fence();
  }

  //----------------------------------------

  struct TagAllocDealloc {};

  KOKKOS_INLINE_FUNCTION
  void operator()(TagAllocDealloc, int i, long& update) const noexcept {
    if (0 == i % fill_stride) {
      const int j = i / fill_stride;

      if (0 == j % 3) {
        for (unsigned k = 0; k < repeat_inner; ++k) {
          const unsigned size_alloc = chunk * (1 + (j % chunk_span));

          pool.deallocate((void*)ptrs(j), size_alloc);

          ptrs(j) = (uintptr_t)pool.allocate(size_alloc);

          if (0 == ptrs(j)) update++;
        }
      }
    }
  }

  bool test_alloc_dealloc() {
    using policy = Kokkos::RangePolicy<ExecSpace, TagAllocDealloc>;

    long error_count = 0;

    Kokkos::parallel_reduce(policy(0, range_iter), *this, error_count);

    return 0 == error_count;
  }
};

int main(int argc, char* argv[]) {
  static const char help_flag[]         = "--help";
  static const char alloc_size_flag[]   = "--alloc_size=";
  static const char super_size_flag[]   = "--super_size=";
  static const char chunk_span_flag[]   = "--chunk_span=";
  static const char fill_stride_flag[]  = "--fill_stride=";
  static const char fill_level_flag[]   = "--fill_level=";
  static const char repeat_outer_flag[] = "--repeat_outer=";
  static const char repeat_inner_flag[] = "--repeat_inner=";

  long total_alloc_size   = 1000000;
  int min_superblock_size = 10000;
  int chunk_span          = 5;
  int fill_stride         = 1;
  int fill_level          = 70;
  int repeat_outer        = 1;
  int repeat_inner        = 1;

  int ask_help = 0;

  for (int i = 1; i < argc; i++) {
    const char* const a = argv[i];

    if (!strncmp(a, help_flag, strlen(help_flag))) ask_help = 1;

    if (!strncmp(a, alloc_size_flag, strlen(alloc_size_flag)))
      total_alloc_size = atol(a + strlen(alloc_size_flag));

    if (!strncmp(a, super_size_flag, strlen(super_size_flag)))
      min_superblock_size = std::stoi(a + strlen(super_size_flag));

    if (!strncmp(a, fill_stride_flag, strlen(fill_stride_flag)))
      fill_stride = std::stoi(a + strlen(fill_stride_flag));

    if (!strncmp(a, fill_level_flag, strlen(fill_level_flag)))
      fill_level = std::stoi(a + strlen(fill_level_flag));

    if (!strncmp(a, chunk_span_flag, strlen(chunk_span_flag)))
      chunk_span = std::stoi(a + strlen(chunk_span_flag));

    if (!strncmp(a, repeat_outer_flag, strlen(repeat_outer_flag)))
      repeat_outer = std::stoi(a + strlen(repeat_outer_flag));

    if (!strncmp(a, repeat_inner_flag, strlen(repeat_inner_flag)))
      repeat_inner = std::stoi(a + strlen(repeat_inner_flag));
  }

  int chunk_span_bytes = 0;
  for (int i = 0; i < chunk_span; ++i) {
    auto chunk_bytes = TestFunctor::chunk * (1 + i);
    if (chunk_bytes < 64) chunk_bytes = 64;
    auto block_bytes_lg2 =
        Kokkos::Impl::integral_power_of_two_that_contains(chunk_bytes);
    auto block_bytes = (1 << block_bytes_lg2);
    chunk_span_bytes += block_bytes;
  }
  auto actual_superblock_bytes_lg2 =
      Kokkos::Impl::integral_power_of_two_that_contains(min_superblock_size);
  auto actual_superblock_bytes = (1 << actual_superblock_bytes_lg2);
  auto superblock_mask         = actual_superblock_bytes - 1;
  auto nsuperblocks =
      (total_alloc_size + superblock_mask) >> actual_superblock_bytes_lg2;
  auto actual_total_bytes = nsuperblocks * actual_superblock_bytes;
  auto bytes_wanted       = (actual_total_bytes * fill_level) / 100;
  auto chunk_spans        = bytes_wanted / chunk_span_bytes;
  auto number_alloc       = int(chunk_spans * chunk_span);

  if (ask_help) {
    std::cout << "command line options:"
              << " " << help_flag << " " << alloc_size_flag << "##"
              << " " << super_size_flag << "##"
              << " " << fill_stride_flag << "##"
              << " " << fill_level_flag << "##"
              << " " << chunk_span_flag << "##"
              << " " << repeat_outer_flag << "##"
              << " " << repeat_inner_flag << "##" << std::endl;
    return 0;
  }

  Kokkos::initialize(argc, argv);

  double sum_fill_time  = 0;
  double sum_cycle_time = 0;
  double sum_both_time  = 0;
  double min_fill_time  = std::numeric_limits<double>::max();
  double min_cycle_time = std::numeric_limits<double>::max();
  double min_both_time  = std::numeric_limits<double>::max();
  // one alloc in fill, alloc/dealloc pair in repeat_inner
  for (int i = 0; i < repeat_outer; ++i) {
    TestFunctor functor(total_alloc_size, min_superblock_size, number_alloc,
                        fill_stride, chunk_span, repeat_inner);

    Kokkos::Impl::Timer timer;

    if (!functor.test_fill()) {
      Kokkos::abort("fill ");
    }

    auto t0 = timer.seconds();

    if (!functor.test_alloc_dealloc()) {
      Kokkos::abort("alloc/dealloc ");
    }

    auto t1              = timer.seconds();
    auto this_fill_time  = t0;
    auto this_cycle_time = t1 - t0;
    auto this_both_time  = t1;
    sum_fill_time += this_fill_time;
    sum_cycle_time += this_cycle_time;
    sum_both_time += this_both_time;
    min_fill_time  = std::min(min_fill_time, this_fill_time);
    min_cycle_time = std::min(min_cycle_time, this_cycle_time);
    min_both_time  = std::min(min_both_time, this_both_time);
  }

  Kokkos::finalize();

  printf(
      "\"mempool: alloc super stride level span inner outer number\" %ld %d %d "
      "%d %d %d %d %d\n",
      total_alloc_size, min_superblock_size, fill_stride, fill_level,
      chunk_span, repeat_inner, repeat_outer, number_alloc);

  auto avg_fill_time  = sum_fill_time / repeat_outer;
  auto avg_cycle_time = sum_cycle_time / repeat_outer;
  auto avg_both_time  = sum_both_time / repeat_outer;

  printf("\"mempool: fill time (min, avg)\" %.8f %.8f\n", min_fill_time,
         avg_fill_time);

  printf("\"mempool: cycle time (min, avg)\" %.8f %.8f\n", min_cycle_time,
         avg_cycle_time);

  printf("\"mempool: test time (min, avg)\" %.8f %.8f\n", min_both_time,
         avg_both_time);

  printf("\"mempool: fill ops per second (max, avg)\" %g %g\n",
         number_alloc / min_fill_time, number_alloc / avg_fill_time);

  printf("\"mempool: cycle ops per second (max, avg)\" %g %g\n",
         (2 * number_alloc * repeat_inner) / min_cycle_time,
         (2 * number_alloc * repeat_inner) / avg_cycle_time);
}