<commit_before>/**
 * @file
 * 
 * cuRAND boilerplate.
 */
#pragma once

#include "numbirch/cuda/cuda.hpp"
#include "numbirch/common/stl.hpp"
#include "numbirch/macro.hpp"

#include <curand_kernel.h>
#include <cassert>

/*
 * Call a cuRAND function and assert success.
 */
#define CURAND_CHECK(call) \
    { \
      curandStatus_t err = call; \
      assert(err == CURAND_STATUS_SUCCESS); \
      if (CUDA_SYNC) { \
        cudaError_t err = cudaStreamSynchronize(stream); \
        assert(err == cudaSuccess); \
      } \
    }


namespace numbirch {
/**
 * @internal
 * 
 * Device pseudorandom number generators for each host thread. The XORWOW
 * generator of cuRAND is used. The MTGP32 generator may be preferred for its
 * longer period (XORWOW: 2^190, MTGP32: 2^11214), but is more difficult to
 * use when implementing rejection sampling algorithms (e.g. custom
 * implementation for gamma variates here), as it can only generate variates
 * for all threads simultaneously. Furthermore, on inspection of
 * curand_uniform.h, it is documented as generating double-precision variates
 * with only 32 random bits, which does not seem ideal.
 */
extern thread_local curandState_t* rngs;

/**
 * @internal
 * 
 * Initialize cuRAND integrations. This should be called during init() by the
 * backend.
 */
void curand_init();

/**
 * @internal
 * 
 * Terminate cuRAND integrations. This should be called during term() by the
 * backend.
 */
void curand_term();

/**
 * @internal
 * 
 * Get the pseudorandom number generator for the current device thread.
 */
NUMBIRCH_DEVICE inline curandState_t* curand_rng() {
  extern __shared__ curandState_t* shared[];
  curandState_t* rngs = shared[0];
  auto x = threadIdx.x + blockIdx.x*blockDim.x;
  auto y = threadIdx.y + blockIdx.y*blockDim.y;
  return &rngs[x + y*gridDim.x*blockDim.x];
}

template<class T>
NUMBIRCH_DEVICE T curand_gamma_generic(curandState_t* state, const T k) {
  /* based on the implementation of std::gamma_distribution in libc++, part of
   * the LLVM project */
  T x;
  if (k == 1) {
    if constexpr (std::is_same_v<T,double>) {
      x = -std::log(curand_uniform_double(state));
    } else {
      x = -std::log(curand_uniform(state));
    }
  } else if (k > 1) {
    T b = k - 1;
    T c = 3*k - T(0.75);
    while (true) {
      T u, v;
      if constexpr (std::is_same_v<T,double>) {
        u = curand_uniform_double(state);
        v = curand_uniform_double(state);
      } else {
        u = curand_uniform(state);
        v = curand_uniform(state);
      }
      T w = u*(1 - u);
      if (w != 0) {
        T y = std::sqrt(c/w)*(u - T(0.5));
        x = b + y;
        if (x >= 0) {
          T z = 64*w*w*w*v*v;
          if (z <= 1 - 2*y*y/x) {
            break;
          }
          if (std::log(z) <= 2*(b*std::log(x/b) - y)) {
            break;
          }
        }
      }
    }
  } else {
    while (true) {
      T u, v;
      if constexpr (std::is_same_v<T,double>) {
        u = curand_uniform_double(state);
        v = curand_uniform_double(state);
      } else {
        u = curand_uniform(state);
        v = curand_uniform(state);
      }
      if (u <= 1 - k) {
        x = std::pow(u, 1/k);
        if (x <= v) {
          break;
        }
      } else {
        T e = -std::log((1 - u)/k);
        x = std::pow(1 - k + k*e, 1/k);
        if (x <= e + v) {
          break;
        }
      }
    }
  }
  return x;
}

NUMBIRCH_DEVICE inline float curand_gamma(curandState_t* state,
    const float k) {
  return curand_gamma_generic<float>(state, k);
}

NUMBIRCH_DEVICE inline double curand_gamma_double(curandState_t* state,
    const double k) {
  return curand_gamma_generic<double>(state, k);
}

template<class T>
NUMBIRCH_DEVICE int curand_binomial_generic(curandState_t* state, const int n,
    const T ρ) {
  /* based on the implementation of std::binomial_distribution in libc++, part
   * of the LLVM project, which cites <NAME>. (1986). A modal method for
   * generating binomial variables. Communication in Statistics- Theory and
   * Methods. 15(3), 805-813 */
  if (n == 0 || ρ == 0) {<commit_msg>Updated gamma simulation to reduce branches for GPU. Added citations.<commit_after>/**
 * @file
 * 
 * cuRAND boilerplate.
 */
#pragma once

#include "numbirch/cuda/cuda.hpp"
#include "numbirch/common/stl.hpp"
#include "numbirch/macro.hpp"

#include <curand_kernel.h>
#include <cassert>

/*
 * Call a cuRAND function and assert success.
 */
#define CURAND_CHECK(call) \
    { \
      curandStatus_t err = call; \
      assert(err == CURAND_STATUS_SUCCESS); \
      if (CUDA_SYNC) { \
        cudaError_t err = cudaStreamSynchronize(stream); \
        assert(err == cudaSuccess); \
      } \
    }


namespace numbirch {
/**
 * @internal
 * 
 * Device pseudorandom number generators for each host thread. The XORWOW
 * generator of cuRAND is used. The MTGP32 generator may be preferred for its
 * longer period (XORWOW: 2^190, MTGP32: 2^11214), but is more difficult to
 * use when implementing rejection sampling algorithms (e.g. custom
 * implementation for gamma variates here), as it can only generate variates
 * for all threads simultaneously. Furthermore, on inspection of
 * curand_uniform.h, it is documented as generating double-precision variates
 * with only 32 random bits, which does not seem ideal.
 */
extern thread_local curandState_t* rngs;

/**
 * @internal
 * 
 * Initialize cuRAND integrations. This should be called during init() by the
 * backend.
 */
void curand_init();

/**
 * @internal
 * 
 * Terminate cuRAND integrations. This should be called during term() by the
 * backend.
 */
void curand_term();

/**
 * @internal
 * 
 * Get the pseudorandom number generator for the current device thread.
 */
NUMBIRCH_DEVICE inline curandState_t* curand_rng() {
  extern __shared__ curandState_t* shared[];
  curandState_t* rngs = shared[0];
  auto x = threadIdx.x + blockIdx.x*blockDim.x;
  auto y = threadIdx.y + blockIdx.y*blockDim.y;
  return &rngs[x + y*gridDim.x*blockDim.x];
}

template<class T>
NUMBIRCH_DEVICE T curand_gamma_generic(curandState_t* state, const T k) {
  /* see Devroye, L. (1986). Non-Uniform Random Variate Generation.
   * Springer-Verlag, New York, online at
   * http://luc.devroye.org/rnbookindex.html. Specific references below. The
   * same algorithm is used for the k > 1 case in libc++ */
  T a = k;

  /* Best's rejection algorithm works only for a >= 1, but we can use the
   * property that, for G ~ Gamma(a + 1) and U ~ Uniform(0, 1), G*pow(U, 1/a)
   * ~ Gamma(a); see Devroye 1986, p. 420 */
  T scale = T(1.0);
  if (a < T(1.0)) {
    T u;
    if constexpr (std::is_same_v<T,double>) {
      u = curand_uniform_double(state);
    } else {
      u = curand_uniform(state);
    }
    scale = std::pow(u, T(1.0)/a);
    a = k + T(1.0);
  }

  /* Best's rejection algorithm; see Devroye 1986, p. 410 */
  T x;
  T b = a - T(1.0);
  T c = T(3.0)*a - T(0.75);
  bool accept = false;
  do {
    T u, v;
    if constexpr (std::is_same_v<T,double>) {
      u = curand_uniform_double(state);
      v = curand_uniform_double(state);
    } else {
      u = curand_uniform(state);
      v = curand_uniform(state);
    }
    T w = u*(T(1.0) - u);
    if (w != T(0.0)) {
      T y = std::sqrt(c/w)*(u - T(0.5));
      x = b + y;
      if (x >= T(0.0)) {
        T z = T(64.0)*w*w*w*v*v;
        accept = (z <= T(1.0) - T(2.0)*y*y/x) ||
            (std::log(z) <= T(2.0)*(b*std::log(x/b) - y));
      }
    }
  } while (!accept);
  return scale*x;
}

NUMBIRCH_DEVICE inline float curand_gamma(curandState_t* state,
    const float k) {
  return curand_gamma_generic<float>(state, k);
}

NUMBIRCH_DEVICE inline double curand_gamma_double(curandState_t* state,
    const double k) {
  return curand_gamma_generic<double>(state, k);
}

template<class T>
NUMBIRCH_DEVICE int curand_binomial_generic(curandState_t* state, const int n,
    const T ρ) {
  /* based on the implementation of std::binomial_distribution in libc++, part
   * of the LLVM project, which cites <NAME>. (1986). A modal method for
   * generating binomial variables. Communication in Statistics- Theory and
   * Methods. 15(3), 805-813 */
  if (n == 0 || ρ == 0) {
{'fraction_non_alphanumeric': 0.13153395212506105, 'fraction_numerical': 0.017342452369320957, 'mean_word_length': 2.782448036951501, 'pattern_counts': {'":': 0, '<': 30, '<?xml version=': 0, '>': 27, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1344387', 'n_tokens_mistral': 3034, 'n_tokens_neox': 2873, 'n_words': 825}