kokkos_init_tester.cpp

#include <Kokkos_Core.hpp>
#include <Kokkos_Random.hpp>

#include <iostream>
#include <string>
#include <stdio.h>

constexpr size_t DEFAULT_BUFFER_SIZE = 8192; // Number of streams initialized per batch
constexpr size_t Gi = 1024 * 1024 * 1024;

// Enumeration of the generators
enum class RNGType {
  SFC64,
  XorShift64,
  XorShift1024,
  Undefined
};

using ExecutionSpace = Kokkos::DefaultExecutionSpace;
using HostExecutionSpace = Kokkos::DefaultHostExecutionSpace;
namespace KP = Kokkos::Profiling;


// PRNG's initialization test function
template <RNGType Rng>
void test_init(const uint64_t seed, const size_t buffer_size) {
  
  // Allocate device view
  Kokkos::View<uint64_t*, ExecutionSpace> d_buffer("device_buffer", buffer_size);
  // Create host mirror. If ExecutionSpace is host (e.g., Serial/OpenMP), 
  // Kokkos just does a reference increment and no memory is duplicated.
  auto h_buffer = Kokkos::create_mirror_view(d_buffer);

  const size_t buffer_size_in_byte = buffer_size * sizeof(uint64_t);
  const size_t nb_buffer_by_GiB = Gi / buffer_size_in_byte; 
  size_t loop_count = 0;
  uint64_t current_offset = 0;

  // Base XorShift64 instance initialized exactly like Kokkos pool does internally.
  // It is safely instantiated here but only advanced and used if a XorShift RNG is selected.
  Kokkos::Random_XorShift64<HostExecutionSpace> base_gen(seed, 0);
  if constexpr (Rng == RNGType::XorShift64 || Rng == RNGType::XorShift1024) {
    for (int i = 0; i < 17; i++) base_gen.rand();
  }

  while(true) {
    // =========================================================================
    // Generation Step (Factorized by Generator Type)
    // =========================================================================
    if constexpr (Rng == RNGType::SFC64) {
      KP::pushRegion("compute:sfc64_pool_instantiation");
      Kokkos::Random_SFC64_Pool<ExecutionSpace> random_pool(seed, current_offset, buffer_size);
      KP::popRegion();

      Kokkos::fence();
      KP::pushRegion("compute:gen_first_number_sfc64");
      Kokkos::parallel_for("sfc64_init_test", buffer_size, KOKKOS_LAMBDA(const size_t i) {
        auto stream_gen = random_pool.get_state(i);
        // for(int j = 0; j < 18; j++)
        //   stream_gen.urand64();
        d_buffer(i) = stream_gen.urand64();
        random_pool.free_state(stream_gen);
      });
      KP::popRegion();
      // std::cout << std::endl;
      // Deep copy to Host. Becomes a no-op if memory spaces are identical.
      deep_copy(h_buffer, d_buffer); 
      current_offset += buffer_size;
    }
    else if constexpr (Rng == RNGType::XorShift64) {
      KP::pushRegion("compute:gen_first_number_xorshift64");
      // Execute sequentially on host to perfectly replicate the pool state initialization
      for (size_t i = 0; i < buffer_size; ++i) {
        int n1 = base_gen.rand(); int n2 = base_gen.rand();
        int n3 = base_gen.rand(); int n4 = base_gen.rand();
        
        uint64_t state = (((static_cast<uint64_t>(n1)) & 0xffff) << 00) |
                         (((static_cast<uint64_t>(n2)) & 0xffff) << 16) |
                         (((static_cast<uint64_t>(n3)) & 0xffff) << 32) |
                         (((static_cast<uint64_t>(n4)) & 0xffff) << 48);

        // Pass the state to a new disposable instance
        Kokkos::Random_XorShift64<HostExecutionSpace> stream_gen(state);
        h_buffer(i) = stream_gen.urand64();
      }
      KP::popRegion();
    } 
    else if constexpr (Rng == RNGType::XorShift1024) {
      KP::pushRegion("compute:gen_first_number_xorshift1024");
      // Execute sequentially on host to perfectly replicate the pool state initialization

      // Local host views
      Kokkos::View<uint64_t*[16], HostExecutionSpace> state("local_state", 1);
      Kokkos::View<int*, HostExecutionSpace> p("local_p", 1);
      
      for (size_t i = 0; i < buffer_size; ++i) {
        for (int j = 0; j < 16; j++) {
          int n1 = base_gen.rand(); int n2 = base_gen.rand();
          int n3 = base_gen.rand(); int n4 = base_gen.rand();
          
          state(0, j) = (((static_cast<uint64_t>(n1)) & 0xffff) << 00) |
                        (((static_cast<uint64_t>(n2)) & 0xffff) << 16) |
                        (((static_cast<uint64_t>(n3)) & 0xffff) << 32) |
                        (((static_cast<uint64_t>(n4)) & 0xffff) << 48);
        }

        // Pass the state to a new disposable instance (0 is the initial pointer index 'p')
        Kokkos::Random_XorShift1024<HostExecutionSpace> stream_gen(state, p(0), 0);
        h_buffer(i) = stream_gen.urand64();
      }
      KP::popRegion();
    }

    // Writing & monitoring step
    // for(int i = 0; i < 100; i++)
    //   std::cout << h_buffer(i) << " ";
    // std::cout << std::endl;
    // return;
    KP::pushRegion("compute:print raw buffer on stdout");
    size_t written = fwrite(h_buffer.data(), sizeof(uint64_t), buffer_size, stdout);
    if (written < buffer_size) {
      fprintf(stderr, "\n[Kokkos Init Tester] Pipe broken or fwrite error.\n");
      break; 
    }
    KP::popRegion();

    loop_count++;
    if (loop_count % nb_buffer_by_GiB == 0) { 
      fprintf(stderr, "\r[Kokkos Init Tester] Streams evaluated : %zu GiB\n", (loop_count * buffer_size_in_byte) / Gi);
      fflush(stderr);
    }
  }
}

int main(int argc, char* argv[])
{
  Kokkos::ScopeGuard guard(argc, argv);

  RNGType rng = RNGType::Undefined; 
  std::string rng_name = "Undefined";
  uint64_t seed = 12345;
  size_t buffer_size = DEFAULT_BUFFER_SIZE;

  // Args parsing
  // Expected: ./kokkos_init_tester [generator] [seed] [buffer_size]
  if (argc >= 2) {
    rng_name = argv[1];
    if (rng_name == "sfc64") rng = RNGType::SFC64;
    else if (rng_name == "xorshift64") rng = RNGType::XorShift64;
    else if (rng_name == "xorshift1024") rng = RNGType::XorShift1024;
  }

  if (argc >= 3) seed = std::stoull(argv[2]);
  if (argc >= 4) buffer_size = std::stoul(argv[3]);

  if (rng == RNGType::Undefined) {
    std::cerr << "[Kokkos Init Tester] No generator provided or unknown generator!\n";
    return -1;
  }

  // Log output 
  std::cerr << "[Kokkos Init Tester] Evaluating Initialization for " << rng_name 
            << " | Datatype: uint64_t"
            << " | Base seed: " << seed 
            << " | Batch of streams: " << buffer_size << std::endl;

  // Execution dispatch
  switch (rng) {
    case RNGType::SFC64:        test_init<RNGType::SFC64>(seed, buffer_size); break;
    case RNGType::XorShift64:   test_init<RNGType::XorShift64>(seed, buffer_size); break;
    case RNGType::XorShift1024: test_init<RNGType::XorShift1024>(seed, buffer_size); break;
    default: break;
  }
  
  return 0;
}