Initialise pinned memory

dmlc · May 6, 2020 · 5ea93b8 · 5ea93b8
1 parent afb55cf
commit 5ea93b8
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 5 deletions.
diff --git a/src/common/device_helpers.cuh b/src/common/device_helpers.cuh
@@ -503,6 +503,15 @@ struct PinnedMemory {
     return xgboost::common::Span<T>(static_cast<T *>(temp_storage), size);
   }
 
+  template <typename T>
+  xgboost::common::Span<T> GetSpan(size_t size, T init) {
+    auto result = this->GetSpan<T>(size);
+    for (auto &e : result) {
+      e = init;
+    }
+    return result;
+  }
+
   void Free() {
     if (temp_storage != nullptr) {
       safe_cuda(cudaFreeHost(temp_storage));

diff --git a/src/tree/updater_gpu_hist.cu b/src/tree/updater_gpu_hist.cu
@@ -631,7 +631,9 @@ struct GPUHistMakerDevice {
     // The set of leaves that can be expanded asynchronously
     auto expand_set = driver.Pop();
     while (!expand_set.empty()) {
-      auto new_candidates = pinned.GetSpan<ExpandEntry>(expand_set.size() * 2);
+      auto new_candidates =
+          pinned.GetSpan<ExpandEntry>(expand_set.size() * 2, ExpandEntry());
+
       for (auto i = 0ull; i < expand_set.size(); i++) {
         auto candidate = expand_set.at(i);
         if (!candidate.IsValid(param, num_leaves)) {
@@ -659,10 +661,6 @@ struct GPUHistMakerDevice {
                                         right_child_nidx, *p_tree,
                                         new_candidates.subspan(i * 2, 2));
           monitor.StopCuda("EvaluateSplits");
-        } else {
-          // Set default
-          new_candidates[i * 2] = ExpandEntry();
-          new_candidates[i * 2 + 1] = ExpandEntry();
         }
       }
       dh::safe_cuda(cudaDeviceSynchronize());