random.cc 2.34 KB
Newer Older
1
2
3
4
5
6
/*!
 *  Copyright (c) 2017 by Contributors
 * \file random.cc
 * \brief Random number generator interfaces
 */

7
8
#include <dgl/array.h>
#include <dgl/random.h>
9
#include <dgl/runtime/packed_func.h>
10
#include <dgl/runtime/parallel_for.h>
11
12
#include <dgl/runtime/registry.h>
#include <dmlc/omp.h>
13

14
15
16
17
#ifdef DGL_USE_CUDA
#include "../runtime/cuda/cuda_common.h"
#endif  // DGL_USE_CUDA

18
19
20
21
22
using namespace dgl::runtime;

namespace dgl {

DGL_REGISTER_GLOBAL("rng._CAPI_SetSeed")
23
24
    .set_body([](DGLArgs args, DGLRetValue *rv) {
      const int seed = args[0];
25

26
27
28
29
30
      runtime::parallel_for(0, omp_get_max_threads(), [&](size_t b, size_t e) {
        for (auto i = b; i < e; ++i) {
          RandomEngine::ThreadLocal()->SetSeed(seed);
        }
      });
31
#ifdef DGL_USE_CUDA
32
33
34
35
36
37
38
39
      if (DeviceAPI::Get(kDGLCUDA)->IsAvailable()) {
        auto *thr_entry = CUDAThreadEntry::ThreadLocal();
        if (!thr_entry->curand_gen) {
          CURAND_CALL(curandCreateGenerator(
              &thr_entry->curand_gen, CURAND_RNG_PSEUDO_DEFAULT));
        }
        CURAND_CALL(curandSetPseudoRandomGeneratorSeed(
            thr_entry->curand_gen, static_cast<uint64_t>(seed)));
40
      }
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
41
#endif  // DGL_USE_CUDA
42
    });
43

44
DGL_REGISTER_GLOBAL("rng._CAPI_Choice")
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
    .set_body([](DGLArgs args, DGLRetValue *rv) {
      const int64_t num = args[0];
      const int64_t population = args[1];
      const NDArray prob = args[2];
      const bool replace = args[3];
      const int bits = args[4];
      CHECK(bits == 32 || bits == 64)
          << "Supported bit widths are 32 and 64, but got " << bits << ".";
      if (aten::IsNullArray(prob)) {
        if (bits == 32) {
          *rv = RandomEngine::ThreadLocal()->UniformChoice<int32_t>(
              num, population, replace);
        } else {
          *rv = RandomEngine::ThreadLocal()->UniformChoice<int64_t>(
              num, population, replace);
        }
61
      } else {
62
63
64
65
66
67
68
69
70
71
72
        if (bits == 32) {
          ATEN_FLOAT_TYPE_SWITCH(prob->dtype, FloatType, "probability", {
            *rv = RandomEngine::ThreadLocal()->Choice<int32_t, FloatType>(
                num, prob, replace);
          });
        } else {
          ATEN_FLOAT_TYPE_SWITCH(prob->dtype, FloatType, "probability", {
            *rv = RandomEngine::ThreadLocal()->Choice<int64_t, FloatType>(
                num, prob, replace);
          });
        }
73
      }
74
    });
75

76
};  // namespace dgl