Support platforms without rand_r

Author: leezu

example/multi_threaded_inference/multi_threaded_inference.cc

@@ -34,6 +34,7 @@
 #include <opencv2/opencv.hpp>
 #include <mxnet/c_predict_api.h>
 #include "mxnet-cpp/MxNetCpp.h"
+#include <random>
 
 const float DEFAULT_MEAN = 117.0;
 
@@ -248,7 +249,9 @@ void run_inference(const std::string& model_name, const std::vector<mxnet::cpp::
   auto func = [&](int num) {
     unsigned next = num;
     if (random_sleep) {
-      int sleep_time = rand_r(&next) % 5;
+      static thread_local std::mt19937 generator;
+      std::uniform_int_distribution<int> distribution(0, 5);
+      int sleep_time = distribution(generator);
       std::this_thread::sleep_for(std::chrono::seconds(sleep_time));
     }
     int num_output = 0;

src/operator/contrib/dgl_graph.cc

@@ -24,6 +24,9 @@
 #include <mxnet/operator_util.h>
 #include <dmlc/logging.h>
 #include <dmlc/optional.h>
+#include <algorithm>
+#include <random>
+
 #include "../elemwise_op_common.h"
 #include "../../imperative/imperative_utils.h"
 #include "../subgraph_op_common.h"
@@ -41,7 +44,9 @@ typedef int64_t dgl_id_t;
  */
 class ArrayHeap {
  public:
-  explicit ArrayHeap(const std::vector<float>& prob) {
+  explicit ArrayHeap(const std::vector<float>& prob, unsigned int seed) {
+    generator_ = std::mt19937(seed);
+    distribution_ = std::uniform_real_distribution<float>(0.0, 1.0);
     vec_size_ = prob.size();
     bit_len_ = ceil(log2(vec_size_));
     limit_ = 1 << bit_len_;
@@ -86,8 +91,8 @@ class ArrayHeap {
   /*
    * Sample from arrayHeap
    */
-  size_t Sample(unsigned int* seed) {
-    float xi = heap_[1] * (rand_r(seed)%100/101.0);
+  size_t Sample() {
+    float xi = heap_[1] * distribution_(generator_);
     int i = 1;
     while (i < limit_) {
       i = i << 1;
@@ -102,10 +107,10 @@ class ArrayHeap {
   /*
    * Sample a vector by given the size n
    */
-  void SampleWithoutReplacement(size_t n, std::vector<size_t>* samples, unsigned int* seed) {
+  void SampleWithoutReplacement(size_t n, std::vector<size_t>* samples) {
     // sample n elements
     for (size_t i = 0; i < n; ++i) {
-      samples->at(i) = this->Sample(seed);
+      samples->at(i) = this->Sample();
       this->Delete(samples->at(i));
     }
   }
@@ -115,6 +120,8 @@ class ArrayHeap {
   int bit_len_;   // bit size
   int limit_;
   std::vector<float> heap_;
+  std::mt19937 generator_;
+  std::uniform_real_distribution<float> distribution_;
 };
 
 struct NeighborSampleParam : public dmlc::Parameter<NeighborSampleParam> {
@@ -402,10 +409,12 @@ static bool CSRNeighborNonUniformSampleType(const nnvm::NodeAttrs& attrs,
 static void RandomSample(size_t set_size,
                          size_t num,
                          std::vector<size_t>* out,
-                         unsigned int* seed) {
+                         unsigned int seed) {
+  std::mt19937 generator(seed);
   std::unordered_set<size_t> sampled_idxs;
+  std::uniform_int_distribution<size_t> distribution(0, set_size - 1);
   while (sampled_idxs.size() < num) {
-    sampled_idxs.insert(rand_r(seed) % set_size);
+    sampled_idxs.insert(distribution(generator));
   }
   out->clear();
   for (auto it = sampled_idxs.begin(); it != sampled_idxs.end(); it++) {
@@ -441,7 +450,7 @@ static void GetUniformSample(const dgl_id_t* val_list,
                              const size_t max_num_neighbor,
                              std::vector<dgl_id_t>* out_ver,
                              std::vector<dgl_id_t>* out_edge,
-                             unsigned int* seed) {
+                             unsigned int seed) {
   // Copy ver_list to output
   if (ver_len <= max_num_neighbor) {
     for (size_t i = 0; i < ver_len; ++i) {
@@ -485,7 +494,7 @@ static void GetNonUniformSample(const float* probability,
                                 const size_t max_num_neighbor,
                                 std::vector<dgl_id_t>* out_ver,
                                 std::vector<dgl_id_t>* out_edge,
-                                unsigned int* seed) {
+                                unsigned int seed) {
   // Copy ver_list to output
   if (ver_len <= max_num_neighbor) {
     for (size_t i = 0; i < ver_len; ++i) {
@@ -500,8 +509,8 @@ static void GetNonUniformSample(const float* probability,
   for (size_t i = 0; i < ver_len; ++i) {
     sp_prob[i] = probability[col_list[i]];
   }
-  ArrayHeap arrayHeap(sp_prob);
-  arrayHeap.SampleWithoutReplacement(max_num_neighbor, &sp_index, seed);
+  ArrayHeap arrayHeap(sp_prob, seed);
+  arrayHeap.SampleWithoutReplacement(max_num_neighbor, &sp_index);
   out_ver->resize(max_num_neighbor);
   out_edge->resize(max_num_neighbor);
   for (size_t i = 0; i < max_num_neighbor; ++i) {
@@ -536,8 +545,8 @@ static void SampleSubgraph(const NDArray &csr,
                            const float* probability,
                            int num_hops,
                            size_t num_neighbor,
-                           size_t max_num_vertices) {
-  unsigned int time_seed = time(nullptr);
+                           size_t max_num_vertices,
+                           unsigned int random_seed) {
   size_t num_seeds = seed_arr.shape().Size();
   CHECK_GE(max_num_vertices, num_seeds);
 
@@ -594,7 +603,7 @@ static void SampleSubgraph(const NDArray &csr,
                        num_neighbor,
                        &tmp_sampled_src_list,
                        &tmp_sampled_edge_list,
-                       &time_seed);
+                       random_seed);
     } else {  // non-uniform-sample
       GetNonUniformSample(probability,
                        val_list + *(indptr + dst_id),
@@ -603,7 +612,7 @@ static void SampleSubgraph(const NDArray &csr,
                        num_neighbor,
                        &tmp_sampled_src_list,
                        &tmp_sampled_edge_list,
-                       &time_seed);
+                       random_seed);
     }
     CHECK_EQ(tmp_sampled_src_list.size(), tmp_sampled_edge_list.size());
     size_t pos = neighbor_list.size();
@@ -720,12 +729,15 @@ static void CSRNeighborUniformSampleComputeExCPU(const nnvm::NodeAttrs& attrs,
                                           const std::vector<NDArray>& inputs,
                                           const std::vector<OpReqType>& req,
                                           const std::vector<NDArray>& outputs) {
-  const NeighborSampleParam& params =
-    nnvm::get<NeighborSampleParam>(attrs.parsed);
+  const NeighborSampleParam& params = nnvm::get<NeighborSampleParam>(attrs.parsed);
 
   int num_subgraphs = inputs.size() - 1;
   CHECK_EQ(outputs.size(), 3 * num_subgraphs);
 
+  mshadow::Stream<cpu> *s = ctx.get_stream<cpu>();
+  mshadow::Random<cpu, unsigned int> *prnd = ctx.requested[0].get_random<cpu, unsigned int>(s);
+  unsigned int seed = prnd->GetRandInt();
+
 #pragma omp parallel for
   for (int i = 0; i < num_subgraphs; i++) {
     SampleSubgraph(inputs[0],                     // graph_csr
@@ -737,7 +749,12 @@ static void CSRNeighborUniformSampleComputeExCPU(const nnvm::NodeAttrs& attrs,
                    nullptr,                       // probability
                    params.num_hops,
                    params.num_neighbor,
-                   params.max_num_vertices);
+                   params.max_num_vertices,
+#if defined(_OPENMP)
+                   seed + omp_get_thread_num());
+#else
+                   seed);
+#endif
   }
 }
 
@@ -798,6 +815,9 @@ of max_num_vertices, and the valid number of vertices is the same as the ones in
 .set_attr<mxnet::FInferShape>("FInferShape", CSRNeighborUniformSampleShape)
 .set_attr<nnvm::FInferType>("FInferType", CSRNeighborUniformSampleType)
 .set_attr<FComputeEx>("FComputeEx<cpu>", CSRNeighborUniformSampleComputeExCPU)
+.set_attr<FResourceRequest>("FResourceRequest", [](const NodeAttrs& attrs) {
+  return std::vector<ResourceRequest>{ResourceRequest::kRandom};
+})
 .add_argument("csr_matrix", "NDArray-or-Symbol", "csr matrix")
 .add_argument("seed_arrays", "NDArray-or-Symbol[]", "seed vertices")
 .set_attr<std::string>("key_var_num_args", "num_args")
@@ -811,14 +831,17 @@ static void CSRNeighborNonUniformSampleComputeExCPU(const nnvm::NodeAttrs& attrs
                                               const std::vector<NDArray>& inputs,
                                               const std::vector<OpReqType>& req,
                                               const std::vector<NDArray>& outputs) {
-  const NeighborSampleParam& params =
-    nnvm::get<NeighborSampleParam>(attrs.parsed);
+  const NeighborSampleParam& params = nnvm::get<NeighborSampleParam>(attrs.parsed);
 
   int num_subgraphs = inputs.size() - 2;
   CHECK_EQ(outputs.size(), 4 * num_subgraphs);
 
   const float* probability = inputs[1].data().dptr<float>();
 
+  mshadow::Stream<cpu> *s = ctx.get_stream<cpu>();
+  mshadow::Random<cpu, unsigned int> *prnd = ctx.requested[0].get_random<cpu, unsigned int>(s);
+  unsigned int seed = prnd->GetRandInt();
+
 #pragma omp parallel for
   for (int i = 0; i < num_subgraphs; i++) {
     float* sub_prob = outputs[i+2*num_subgraphs].data().dptr<float>();
@@ -831,7 +854,12 @@ static void CSRNeighborNonUniformSampleComputeExCPU(const nnvm::NodeAttrs& attrs
                    probability,
                    params.num_hops,
                    params.num_neighbor,
-                   params.max_num_vertices);
+                   params.max_num_vertices,
+#if defined(_OPENMP)
+                   seed + omp_get_thread_num());
+#else
+                   seed);
+#endif
   }
 }
 
@@ -897,6 +925,9 @@ of max_num_vertices, and the valid number of vertices is the same as the ones in
 .set_attr<mxnet::FInferShape>("FInferShape", CSRNeighborNonUniformSampleShape)
 .set_attr<nnvm::FInferType>("FInferType", CSRNeighborNonUniformSampleType)
 .set_attr<FComputeEx>("FComputeEx<cpu>", CSRNeighborNonUniformSampleComputeExCPU)
+.set_attr<FResourceRequest>("FResourceRequest", [](const NodeAttrs& attrs) {
+  return std::vector<ResourceRequest>{ResourceRequest::kRandom};
+})
 .add_argument("csr_matrix", "NDArray-or-Symbol", "csr matrix")
 .add_argument("probability", "NDArray-or-Symbol", "probability vector")
 .add_argument("seed_arrays", "NDArray-or-Symbol[]", "seed vertices")

src/operator/nn/mkldnn/mkldnn_rnn.cc

@@ -995,7 +995,7 @@ void MKLDNNRnnOp::Forward(const OpContext &ctx,
                           const std::vector<NDArray> &inputs,
                           const std::vector<OpReqType> &req,
                           const std::vector<NDArray> &outputs) {
-  TmpMemMgr::Get()->Init(ctx.requested[0]);
+  TmpMemMgr::Get()->Init(ctx.requested[1]);
   // In the `autograd.record()` context, RNNOp is required to run into
   // forward_training mode.
   const bool is_training = (ctx.is_train || ctx.need_grad);
@@ -1132,7 +1132,7 @@ void MKLDNNRnnOp::Backward(const OpContext& ctx,
                            const std::vector<OpReqType>& req,
                            const std::vector<NDArray>& outputs) {
   using tag = mkldnn::memory::format_tag;
-  TmpMemMgr::Get()->Init(ctx.requested[0]);
+  TmpMemMgr::Get()->Init(ctx.requested[1]);
   const RNNParam& default_param = full_param_.default_param;
   const int data_dtype = inputs[rnn_enum::kData].dtype();
   const int w_dtype = inputs[rnn_enum::kParams].dtype();

src/operator/rnn-inl.h

@@ -31,6 +31,7 @@
 #include <mxnet/operator.h>
 #include <mxnet/storage.h>
 #include <algorithm>
+#include <random>
 #include <map>
 #include <vector>
 #include <string>
@@ -293,23 +294,24 @@ void RNNForwardTraining(DType* ws,
                         DType* hy_ptr,
                         DType* cy_ptr,
                         const float dropout,
-                        int mode) {
+                        int mode,
+                        std::mt19937 &rnd_engine) {  // NOLINT(runtime/references)
   switch (mode) {
     case rnn_enum::kLstm:
       LstmForwardTraining<DType>(ws, rs, state_outputs, num_layers, direction, seq_length,
                                  batch_size, input_size, state_size, x_ptr, hx_ptr, cx_ptr,
-                                 w_ptr, b_ptr, y_ptr, hy_ptr, cy_ptr, dropout);
+                                 w_ptr, b_ptr, y_ptr, hy_ptr, cy_ptr, dropout, rnd_engine);
       break;
     case rnn_enum::kGru:
       GruForwardTraining<DType>(ws, rs, state_outputs, num_layers, direction, seq_length,
                                 batch_size, input_size, state_size, x_ptr, hx_ptr,
-                                w_ptr, y_ptr, hy_ptr, dropout);
+                                w_ptr, y_ptr, hy_ptr, dropout, rnd_engine);
       break;
     case rnn_enum::kRnnTanh:
     case rnn_enum::kRnnRelu:
       VanillaRNNForwardTraining<DType>(ws, rs, state_outputs, num_layers, direction, seq_length,
                                        batch_size, input_size, state_size, x_ptr, hx_ptr,
-                                       w_ptr, y_ptr, hy_ptr, dropout, mode);
+                                       w_ptr, y_ptr, hy_ptr, dropout, mode, rnd_engine);
       break;
     default:
       LOG(FATAL) << "unknown RNN mode " << mode;
@@ -842,7 +844,8 @@ class RNNOp {
     }
 #endif  // MXNET_USE_CUDNN == 1 && defined(__CUDACC__)
 
-    if (ctx_.dev_type == kCPU) {
+#if !defined(__CUDACC__)  // cuda doesn't support C++17
+    if constexpr (std::is_same<xpu, cpu>::value) {
       int projection_size = 0;
       if (param_.projection_size.has_value()) {
         projection_size = param_.projection_size.value();
@@ -860,6 +863,9 @@ class RNNOp {
       DType* work_cpu_space = static_cast<DType*>(temp_cpu_space_.data().dptr_);
 
       if (ctx.is_train || ctx.need_grad) {
+        mshadow::Random<cpu, unsigned> *prnd = ctx.requested[0].get_random<xpu, unsigned int>(s);
+        std::mt19937 &rnd_engine = prnd->GetRndEngine();
+
         // allocate reserve space
         if (param_.projection_size.has_value()) {
           LOG(FATAL) << "No training support for LSTM with projection on CPU currently.";
@@ -894,7 +900,8 @@ class RNNOp {
                                   hy_ptr,
                                   cy_ptr,
                                   param_.p,
-                                  param_.mode);
+                                  param_.mode,
+                                  rnd_engine);
       } else {
         RNNForwardInference<DType>(work_cpu_space,
                                    param_.state_outputs,
@@ -916,6 +923,7 @@ class RNNOp {
                                    param_.mode);
       }
     }
+#endif
   }
 
   void Backward(const OpContext &ctx,

src/operator/rnn.cc

@@ -184,6 +184,7 @@ static std::vector<ResourceRequest> RNNResourceEx(const NodeAttrs& attrs, const
     }
 #endif
   } else {
+    request.emplace_back(ResourceRequest::kRandom);
 #if MXNET_USE_MKLDNN == 1
     request.emplace_back(ResourceRequest::kTempSpace);
 #endif