From 5e382339e2ece2be5f24a1c9038e111fdce81e11 Mon Sep 17 00:00:00 2001 From: Tianyu Liu Date: Fri, 15 Nov 2024 11:39:19 -0500 Subject: [PATCH] Update --- cpp/include/kvikio/detail/io_operation.hpp | 29 +++++ cpp/include/kvikio/file_handle.hpp | 119 ++++++++++----------- cpp/include/kvikio/posix_io.hpp | 10 +- 3 files changed, 87 insertions(+), 71 deletions(-) create mode 100644 cpp/include/kvikio/detail/io_operation.hpp diff --git a/cpp/include/kvikio/detail/io_operation.hpp b/cpp/include/kvikio/detail/io_operation.hpp new file mode 100644 index 0000000000..b6caef8260 --- /dev/null +++ b/cpp/include/kvikio/detail/io_operation.hpp @@ -0,0 +1,29 @@ +/* + * Copyright (c) 2024, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include + +namespace kvikio::detail { + +/** + * @brief Type of the IO operation. + */ +enum class IOOperationType : uint8_t { + READ, + WRITE, +}; +} // namespace kvikio::detail diff --git a/cpp/include/kvikio/file_handle.hpp b/cpp/include/kvikio/file_handle.hpp index c434b16179..10e15cd999 100644 --- a/cpp/include/kvikio/file_handle.hpp +++ b/cpp/include/kvikio/file_handle.hpp @@ -27,6 +27,7 @@ #include #include #include +#include #include #include #include @@ -51,6 +52,58 @@ class FileHandle { mutable std::size_t _nbytes{0}; // The size of the underlying file, zero means unknown. CUfileHandle_t _handle{}; + template + void handle_async_io(void* devPtr_base, + std::size_t* size_p, + off_t* file_offset_p, + off_t* devPtr_offset_p, + ssize_t* bytes_done_p, + CUstream stream) + { + auto posix_fallback = [&] { + CUDA_DRIVER_TRY(cudaAPI::instance().StreamSynchronize(stream)); + if constexpr (Operation == detail::IOOperationType::READ) { + *bytes_done_p = + static_cast(read(devPtr_base, *size_p, *file_offset_p, *devPtr_offset_p)); + } else { + *bytes_done_p = + static_cast(write(devPtr_base, *size_p, *file_offset_p, *devPtr_offset_p)); + } + }; + + if (defaults::can_compat_mode_reduce_to_off(_compat_mode)) { + if (!kvikio::is_batch_and_stream_available()) { + if (_compat_mode == CompatMode::AUTO) { + posix_fallback(); + return; + } + throw std::runtime_error("Missing cuFile batch or stream library symbol."); + } + + // When checking for availability, we also check if cuFile's config file exist. This is + // because even when the stream API is available, it doesn't work if no config file exist. + if (config_path().empty()) { + if (_compat_mode == CompatMode::AUTO) { + posix_fallback(); + return; + } + throw std::runtime_error("Missing cuFile configuration file."); + } + + if constexpr (Operation == detail::IOOperationType::READ) { + CUFILE_TRY(cuFileAPI::instance().ReadAsync( + _handle, devPtr_base, size_p, file_offset_p, devPtr_offset_p, bytes_done_p, stream)); + } else { + CUFILE_TRY(cuFileAPI::instance().WriteAsync( + _handle, devPtr_base, size_p, file_offset_p, devPtr_offset_p, bytes_done_p, stream)); + } + + return; + } + + posix_fallback(); + } + public: static constexpr mode_t m644 = S_IRUSR | S_IWUSR | S_IRGRP | S_IWGRP | S_IROTH; FileHandle() noexcept = default; @@ -473,37 +526,8 @@ class FileHandle { ssize_t* bytes_read_p, CUstream stream) { - auto posix_fallback = [&] { - CUDA_DRIVER_TRY(cudaAPI::instance().StreamSynchronize(stream)); - *bytes_read_p = - static_cast(read(devPtr_base, *size_p, *file_offset_p, *devPtr_offset_p)); - }; - - if (defaults::can_compat_mode_reduce_to_off(_compat_mode)) { - if (!kvikio::is_batch_and_stream_available()) { - if (_compat_mode == CompatMode::AUTO) { - posix_fallback(); - return; - } - throw std::runtime_error("Missing cuFile batch or stream library symbol."); - } - - // When checking for availability, we also check if cuFile's config file exist. This is - // because even when the stream API is available, it doesn't work if no config file exist. - if (config_path().empty()) { - if (_compat_mode == CompatMode::AUTO) { - posix_fallback(); - return; - } - throw std::runtime_error("Missing cuFile configuration file."); - } - - CUFILE_TRY(cuFileAPI::instance().ReadAsync( - _handle, devPtr_base, size_p, file_offset_p, devPtr_offset_p, bytes_read_p, stream)); - return; - } - - posix_fallback(); + handle_async_io( + devPtr_base, size_p, file_offset_p, devPtr_offset_p, bytes_read_p, stream); } /** @@ -586,37 +610,8 @@ class FileHandle { ssize_t* bytes_written_p, CUstream stream) { - auto posix_fallback = [&] { - CUDA_DRIVER_TRY(cudaAPI::instance().StreamSynchronize(stream)); - *bytes_written_p = - static_cast(write(devPtr_base, *size_p, *file_offset_p, *devPtr_offset_p)); - }; - - if (defaults::can_compat_mode_reduce_to_off(_compat_mode)) { - if (!kvikio::is_batch_and_stream_available()) { - if (_compat_mode == CompatMode::AUTO) { - posix_fallback(); - return; - } - throw std::runtime_error("Missing cuFile batch or stream library symbol."); - } - - // When checking for availability, we also check if cuFile's config file exist. This is - // because even when the stream API is available, it doesn't work if no config file exist. - if (config_path().empty()) { - if (_compat_mode == CompatMode::AUTO) { - posix_fallback(); - return; - } - throw std::runtime_error("Missing cuFile configuration file."); - } - - CUFILE_TRY(cuFileAPI::instance().WriteAsync( - _handle, devPtr_base, size_p, file_offset_p, devPtr_offset_p, bytes_written_p, stream)); - return; - } - - posix_fallback(); + handle_async_io( + devPtr_base, size_p, file_offset_p, devPtr_offset_p, bytes_written_p, stream); } /** diff --git a/cpp/include/kvikio/posix_io.hpp b/cpp/include/kvikio/posix_io.hpp index 4327a301ec..5bf59ccaf1 100644 --- a/cpp/include/kvikio/posix_io.hpp +++ b/cpp/include/kvikio/posix_io.hpp @@ -22,20 +22,12 @@ #include #include +#include #include #include #include namespace kvikio::detail { - -/** - * @brief Type of the IO operation. - */ -enum class IOOperationType : uint8_t { - READ, ///< POSIX read. - WRITE, ///< POSIX write. -}; - /** * @brief Specifies whether all requested bytes are to be processed or not. */