From 794fd6a46b419d557327eb13aeea6b01e1e9105d Mon Sep 17 00:00:00 2001 From: Jiaming Yuan Date: Thu, 25 Mar 2021 09:58:09 +0800 Subject: [PATCH] Support v3 cuda array interface. (#6776) --- src/data/array_interface.cu | 21 +++++++++++ src/data/array_interface.h | 49 ++++++++++++++++++-------- tests/cpp/data/test_array_interface.cu | 42 ++++++++++++++++++++++ 3 files changed, 98 insertions(+), 14 deletions(-) create mode 100644 src/data/array_interface.cu create mode 100644 tests/cpp/data/test_array_interface.cu diff --git a/src/data/array_interface.cu b/src/data/array_interface.cu new file mode 100644 index 000000000..def4de195 --- /dev/null +++ b/src/data/array_interface.cu @@ -0,0 +1,21 @@ +/*! + * Copyright 2021 by Contributors + */ +#include "array_interface.h" +#include "../common/common.h" + +namespace xgboost { +void ArrayInterfaceHandler::SyncCudaStream(int64_t stream) { + switch (stream) { + case 0: + LOG(FATAL) << "Invalid stream ID in array interface: " << stream; + case 1: + // default legacy stream + break; + case 2: + // default per-thread stream + default: + dh::safe_cuda(cudaStreamSynchronize(reinterpret_cast(stream))); + } +} +} // namespace xgboost diff --git a/src/data/array_interface.h b/src/data/array_interface.h index e88c6007f..71db92b6a 100644 --- a/src/data/array_interface.h +++ b/src/data/array_interface.h @@ -18,6 +18,7 @@ #include "xgboost/logging.h" #include "xgboost/span.h" #include "../common/bitfield.h" +#include "../common/common.h" namespace xgboost { // Common errors in parsing columnar format. @@ -41,7 +42,7 @@ struct ArrayInterfaceErrors { return str.c_str(); } static char const* Version() { - return "Only version 1 and 2 of `__cuda_array_interface__' are supported."; + return "Only version <= 3 of `__cuda_array_interface__' are supported."; } static char const* OfType(std::string const& type) { static std::string str; @@ -119,9 +120,18 @@ class ArrayInterfaceHandler { } static void Validate(std::map const& array) { - if (array.find("version") == array.cend()) { + auto version_it = array.find("version"); + if (version_it == array.cend()) { LOG(FATAL) << "Missing `version' field for array interface"; } + auto stream_it = array.find("stream"); + if (stream_it != array.cend() && !IsA(stream_it->second)) { + // is cuda, check the version. + if (get(version_it->second) > 3) { + LOG(FATAL) << ArrayInterfaceErrors::Version(); + } + } + if (array.find("typestr") == array.cend()) { LOG(FATAL) << "Missing `typestr' field for array interface"; } @@ -233,25 +243,31 @@ class ArrayInterfaceHandler { } return p_data; } + + static void SyncCudaStream(int64_t stream); }; +#if !defined(XGBOOST_USE_CUDA) +inline void ArrayInterfaceHandler::SyncCudaStream(int64_t stream) { + common::AssertGPUSupport(); +} +#endif // !defined(XGBOOST_USE_CUDA) + // A view over __array_interface__ class ArrayInterface { - void Initialize(std::map const &column, + void Initialize(std::map const &array, bool allow_mask = true) { - ArrayInterfaceHandler::Validate(column); - auto typestr = get(column.at("typestr")); + ArrayInterfaceHandler::Validate(array); + auto typestr = get(array.at("typestr")); this->AssignType(StringView{typestr}); - auto shape = ArrayInterfaceHandler::ExtractShape(column); - num_rows = shape.first; - num_cols = shape.second; - - data = ArrayInterfaceHandler::ExtractData(column, StringView{typestr}, shape); + std::tie(num_rows, num_cols) = ArrayInterfaceHandler::ExtractShape(array); + data = ArrayInterfaceHandler::ExtractData( + array, StringView{typestr}, std::make_pair(num_rows, num_cols)); if (allow_mask) { common::Span s_mask; - size_t n_bits = ArrayInterfaceHandler::ExtractMask(column, &s_mask); + size_t n_bits = ArrayInterfaceHandler::ExtractMask(array, &s_mask); valid = RBitField8(s_mask); @@ -261,12 +277,18 @@ class ArrayInterface { << "XGBoost doesn't support internal broadcasting."; } } else { - CHECK(column.find("mask") == column.cend()) + CHECK(array.find("mask") == array.cend()) << "Masked array is not yet supported."; } - ArrayInterfaceHandler::ExtractStride(column, strides, num_rows, num_cols, + ArrayInterfaceHandler::ExtractStride(array, strides, num_rows, num_cols, typestr[2] - '0'); + + auto stream_it = array.find("stream"); + if (stream_it != array.cend() && !IsA(stream_it->second)) { + int64_t stream = get(stream_it->second); + ArrayInterfaceHandler::SyncCudaStream(stream); + } } public: @@ -377,7 +399,6 @@ class ArrayInterface { bst_feature_t num_cols; size_t strides[2]{0, 0}; void* data; - Type type; }; diff --git a/tests/cpp/data/test_array_interface.cu b/tests/cpp/data/test_array_interface.cu new file mode 100644 index 000000000..75923e74b --- /dev/null +++ b/tests/cpp/data/test_array_interface.cu @@ -0,0 +1,42 @@ +/*! + * Copyright 2021 by Contributors + */ +#include +#include +#include "../helpers.h" +#include "../../../src/data/array_interface.h" + +namespace xgboost { + +__global__ void SleepForTest(uint64_t *out, uint64_t duration) { + auto start = clock64(); + auto t = 0; + while (t < duration) { + t = clock64() - start; + } + out[0] = t; +} + +TEST(ArrayInterface, Stream) { + size_t constexpr kRows = 10, kCols = 10; + HostDeviceVector storage; + auto arr_str = RandomDataGenerator{kRows, kCols, 0}.GenerateArrayInterface(&storage); + + cudaStream_t stream; + cudaStreamCreate(&stream); + + auto j_arr =Json::Load(StringView{arr_str}); + j_arr["stream"] = Integer(reinterpret_cast(stream)); + Json::Dump(j_arr, &arr_str); + + dh::caching_device_vector out(1, 0); + uint64_t dur = 1e9; + dh::LaunchKernel{1, 1, 0, stream}(SleepForTest, out.data().get(), dur); + ArrayInterface arr(arr_str); + + auto t = out[0]; + CHECK_GE(t, dur); + + cudaStreamDestroy(stream); +} +} // namespace xgboost