kokkos_random_generator.cpp

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

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

constexpr size_t DEFAULT_BUFFER_SIZE = 8192; // * sizeof(uint64_t) = 64 KiB in total
constexpr size_t Gi = 1024 * 1024 * 1024;

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

// Enumeration of the execution modes
enum class ExecMode {
  Basic,
  Mono
};

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

using HostExecutionSpace = Kokkos::DefaultHostExecutionSpace;
using HostMemorySpace = HostExecutionSpace::memory_space;
  
namespace KP = Kokkos::Profiling;

template <typename RandomPoolType, typename DataType>
void execute_prng(ExecMode mode, const uint64_t seed, const size_t buffer_size) {
  static_assert(std::is_same_v<DataType, uint64_t> || std::is_same_v<DataType, uint32_t>,
                "execute_prng only supports uint64_t and uint32_t");

  KP::pushRegion("init:pool_seeding"); 
  RandomPoolType random_pool(ExecutionSpace{}, seed);
  KP::popRegion();

  KP::pushRegion("init:buffers_allocation");
  Kokkos::View<DataType*, ExecutionSpace> d_buffer("device_buffer", buffer_size);
  Kokkos::View<DataType*, HostExecutionSpace> h_buffer("host_buffer", buffer_size);
  KP::popRegion();

  const size_t buffer_size_in_byte = buffer_size * sizeof(DataType);
  const size_t nb_buffer_by_GiB = Gi / buffer_size_in_byte; 
  
  KP::pushRegion("compute:while_loop");
  size_t loop_count = 0;
  
  while(true) {
    // Choice on the execution mode
    if (mode == ExecMode::Basic) {
      Kokkos::parallel_for("gen numbers from multiple streams", buffer_size, KOKKOS_LAMBDA(const size_t i) {
        auto generator = random_pool.get_state();
        
        DataType rnd_val;
        if constexpr (std::is_same_v<DataType, uint64_t>) {
          rnd_val = generator.urand64();
        } else {
          rnd_val = generator.urand();
        }
        random_pool.free_state(generator);

        d_buffer(i) = rnd_val;
      });
    } else { // ExecMode::Mono
      Kokkos::parallel_for("gen numbers from a mono stream (0)", 1, KOKKOS_LAMBDA(const int /*i*/) {
        auto generator = random_pool.get_state(0);
        for(size_t j = 0; j < buffer_size; j++) {
          DataType rnd_val;
          if constexpr (std::is_same_v<DataType, uint64_t>) {
            rnd_val = generator.urand64();
          } else {
            rnd_val = generator.urand();
          }
          d_buffer(j) = rnd_val;
        }
        random_pool.free_state(generator);
      });
    }

    deep_copy(h_buffer, d_buffer); // Copy on host to output on stdout

    KP::pushRegion("compute:print raw buffer on stdout from Host");
    size_t written = fwrite(h_buffer.data(), sizeof(DataType), buffer_size, stdout);
    
    // Error check
    if (written < buffer_size) {
      fprintf(stderr, "\n[Kokkos] Pipe broken or fwrite error.\n");
      break; 
    }

    // Control display also on stderr to not mess up stdout used by PractRand
    loop_count++;
    if (loop_count % nb_buffer_by_GiB == 0) { // print GiB by GiB
      fprintf(stderr, "\r[Kokkos] Data sent : %zu GiB\n",
              (loop_count * buffer_size_in_byte) / Gi);
      fflush(stderr);
    }
    KP::popRegion();
  }
  KP::popRegion();
};

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

  // Default values
  RNGType rng = RNGType::Undefined; // No default the generator should be specified 
  std::string rng_name = "Undefined";
  
  std::string type_to_produce = "uint64";
  uint64_t seed = 12345;
  size_t buffer_size = DEFAULT_BUFFER_SIZE;
  ExecMode mode = ExecMode::Basic;

  // Args parsing
  // Generator choice
  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;
  }

  // Type to produce
  if (argc >= 3) {
    type_to_produce = argv[2];
    if (type_to_produce != "uint32" && type_to_produce != "uint64") {
      std::cerr << "[Kokkos] Unsupported datatype! Only \"uint64\" and \"uint32\" are supported.\n";
      return -1;
    }
  }

  // Seed to use
  if (argc >= 4) seed = std::stoull(argv[3]);
  
  // Size of the buffer to use to output numbers
  if (argc >= 5) buffer_size = std::stoul(argv[4]);

  // Choose the execution mode (mono stream vs multiple stream)
  if (argc >= 6) {
    std::string arg5 = argv[5];
    if (arg5 == "mono") mode = ExecMode::Mono;
    else if (arg5 == "basic") mode = ExecMode::Basic;
    else {
      std::cerr << "[Kokkos] Unknown execution mode! Only \"basic\" and \"mono\" are supported.\n";
      return -2;
    }
  }

  // Assert that the generator has been provided by the user
  if (rng == RNGType::Undefined) {
    std::cerr << "[Kokkos] No generator provided or unknown generator!\n";
    return -3;
  }

  // Log output
  std::cerr << "[Kokkos] Generator " << rng_name 
            << " producing " << type_to_produce 
            << " with seed: " << seed 
            << " in batch of: " << buffer_size << " numbers"
            << " from " << ((mode == ExecMode::Mono)? "a single PRNG stream" : "multiple PRNG streams") << std::endl;

  // Execution dispatch
  if (type_to_produce == "uint64") {
    switch (rng) {
      case RNGType::SFC64:        execute_prng<Kokkos::Random_SFC64_Pool<>, uint64_t>(mode, seed, buffer_size); break;
      case RNGType::XorShift64:   execute_prng<Kokkos::Random_XorShift64_Pool<>, uint64_t>(mode, seed, buffer_size); break;
      case RNGType::XorShift1024: execute_prng<Kokkos::Random_XorShift1024_Pool<>, uint64_t>(mode, seed, buffer_size); break;
      default: break;
    }
  } else { // uint32_t
    switch (rng) {
      case RNGType::SFC64:        execute_prng<Kokkos::Random_SFC64_Pool<>, uint32_t>(mode, seed, buffer_size); break;
      case RNGType::XorShift64:   execute_prng<Kokkos::Random_XorShift64_Pool<>, uint32_t>(mode, seed, buffer_size); break;
      case RNGType::XorShift1024: execute_prng<Kokkos::Random_XorShift1024_Pool<>, uint32_t>(mode, seed, buffer_size); break;
      default: break;
    }
  }
  
  return 0;
}