Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Wait for data CUDA stream instead of sync. #9144

Merged
merged 2 commits into from
May 9, 2023
Merged

Wait for data CUDA stream instead of sync. #9144

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7f50f9d
Wait for data CUDA stream instead of sync.
trivialfis May 8, 2023
c5160d9
Update src/data/array_interface.cu
trivialfis May 9, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 5 additions & 7 deletions src/common/device_helpers.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1352,14 +1352,12 @@ class CUDAStream {
cudaStream_t stream_;

public:
CUDAStream() {
dh::safe_cuda(cudaStreamCreateWithFlags(&stream_, cudaStreamNonBlocking));
}
~CUDAStream() {
dh::safe_cuda(cudaStreamDestroy(stream_));
}
CUDAStream() { dh::safe_cuda(cudaStreamCreateWithFlags(&stream_, cudaStreamNonBlocking)); }
~CUDAStream() { dh::safe_cuda(cudaStreamDestroy(stream_)); }

[[nodiscard]] CUDAStreamView View() const { return CUDAStreamView{stream_}; }
[[nodiscard]] cudaStream_t Handle() const { return stream_; }

CUDAStreamView View() const { return CUDAStreamView{stream_}; }
void Sync() { this->View().Sync(); }
};

Expand Down
17 changes: 12 additions & 5 deletions src/data/array_interface.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,11 +1,15 @@
/*!
* Copyright 2021 by Contributors
/**
* Copyright 2021-2023, XGBoost Contributors
*/
#include <cinttypes> // for int64_t
trivialfis marked this conversation as resolved.
Show resolved Hide resolved

#include "../common/common.h"
#include "../common/device_helpers.cuh" // for DefaultStream, CUDAEvent
#include "array_interface.h"
#include "xgboost/logging.h"

namespace xgboost {
void ArrayInterfaceHandler::SyncCudaStream(int64_t stream) {
void ArrayInterfaceHandler::SyncCudaStream(std::int64_t stream) {
switch (stream) {
case 0:
/**
Expand All @@ -22,8 +26,11 @@ void ArrayInterfaceHandler::SyncCudaStream(int64_t stream) {
break;
case 2:
// default per-thread stream
default:
dh::safe_cuda(cudaStreamSynchronize(reinterpret_cast<cudaStream_t>(stream)));
default: {
dh::CUDAEvent e;
e.Record(dh::CUDAStreamView{reinterpret_cast<cudaStream_t>(stream)});
dh::DefaultStream().Wait(e);
}
}
}

Expand Down
17 changes: 7 additions & 10 deletions tests/cpp/data/test_array_interface.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*!
* Copyright 2021 by Contributors
/**
* Copyright 2021-2023, XGBoost Contributors
*/
#include <gtest/gtest.h>
#include <xgboost/host_device_vector.h>
Expand All @@ -22,22 +22,19 @@ TEST(ArrayInterface, Stream) {
HostDeviceVector<float> storage;
auto arr_str = RandomDataGenerator{kRows, kCols, 0}.GenerateArrayInterface(&storage);

cudaStream_t stream;
cudaStreamCreate(&stream);
dh::CUDAStream stream;

auto j_arr =Json::Load(StringView{arr_str});
j_arr["stream"] = Integer(reinterpret_cast<int64_t>(stream));
auto j_arr = Json::Load(StringView{arr_str});
j_arr["stream"] = Integer(reinterpret_cast<int64_t>(stream.Handle()));
Json::Dump(j_arr, &arr_str);

dh::caching_device_vector<uint64_t> out(1, 0);
uint64_t dur = 1e9;
dh::LaunchKernel{1, 1, 0, stream}(SleepForTest, out.data().get(), dur);
std::uint64_t dur = 1e9;
dh::LaunchKernel{1, 1, 0, stream.View()}(SleepForTest, out.data().get(), dur);
ArrayInterface<2> arr(arr_str);

auto t = out[0];
CHECK_GE(t, dur);

cudaStreamDestroy(stream);
}

TEST(ArrayInterface, Ptr) {
Expand Down