Skip to content

Commit 3f92970

Browse files
authored
[backport] Fix CUDA async stream. (dmlc#8380) (dmlc#8392)
1 parent e17f701 commit 3f92970

File tree

1 file changed

+4
-0
lines changed

1 file changed

+4
-0
lines changed

src/objective/adaptive.cu

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,10 @@ void EncodeTreeLeafDevice(Context const* ctx, common::Span<bst_node_t const> pos
6767
auto pinned = pinned_pool.GetSpan<char>(sizeof(size_t) + sizeof(bst_node_t));
6868
dh::CUDAStream copy_stream;
6969
size_t* h_num_runs = reinterpret_cast<size_t*>(pinned.subspan(0, sizeof(size_t)).data());
70+
71+
dh::CUDAEvent e;
72+
e.Record(dh::DefaultStream());
73+
copy_stream.View().Wait(e);
7074
// flag for whether there's ignored position
7175
bst_node_t* h_first_unique =
7276
reinterpret_cast<bst_node_t*>(pinned.subspan(sizeof(size_t), sizeof(bst_node_t)).data());

0 commit comments

Comments
 (0)