From c2846fb8580a2a6ddd29808bea470623b71feb19 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Mon, 14 Feb 2022 12:32:12 -0800 Subject: [PATCH] Fix incorrect slicing of GDS read/write calls (#10274) Issue happens when the read/write size is a multiple of the maximum slice size. It this case, size of the last slice is computed as `0`, instead of `max_slice_size`: `(t == n_slices - 1) ? size % max_slice_bytes : max_slice_bytes` This PR reimplements this part of code and adds unit tests. Authors: - Vukasin Milovanovic (https://github.com/vuule) Approvers: - Mike Wilson (https://github.com/hyperbolic2346) - Devavret Makkar (https://github.com/devavret) URL: https://github.com/rapidsai/cudf/pull/10274 --- cpp/src/io/utilities/file_io_utilities.cpp | 36 ++++++++++------- cpp/src/io/utilities/file_io_utilities.hpp | 17 +++++++- cpp/tests/CMakeLists.txt | 1 + cpp/tests/io/file_io_test.cpp | 46 ++++++++++++++++++++++ 4 files changed, 85 insertions(+), 15 deletions(-) create mode 100644 cpp/tests/io/file_io_test.cpp diff --git a/cpp/src/io/utilities/file_io_utilities.cpp b/cpp/src/io/utilities/file_io_utilities.cpp index dabe992d959..e2893a2e881 100644 --- a/cpp/src/io/utilities/file_io_utilities.cpp +++ b/cpp/src/io/utilities/file_io_utilities.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021, NVIDIA CORPORATION. + * Copyright (c) 2021-2022, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -194,20 +194,13 @@ template > make_sliced_tasks( F function, DataT* ptr, size_t offset, size_t size, cudf::detail::thread_pool& pool) { + constexpr size_t default_max_slice_size = 4 * 1024 * 1024; + static auto const max_slice_size = getenv_or("LIBCUDF_CUFILE_SLICE_SIZE", default_max_slice_size); + auto const slices = make_file_io_slices(size, max_slice_size); std::vector> slice_tasks; - constexpr size_t default_max_slice_bytes = 4 * 1024 * 1024; - static auto const max_slice_bytes = - getenv_or("LIBCUDF_CUFILE_SLICE_SIZE", default_max_slice_bytes); - size_t const n_slices = util::div_rounding_up_safe(size, max_slice_bytes); - size_t slice_offset = 0; - for (size_t t = 0; t < n_slices; ++t) { - DataT* ptr_slice = ptr + slice_offset; - - size_t const slice_size = (t == n_slices - 1) ? size % max_slice_bytes : max_slice_bytes; - slice_tasks.push_back(pool.submit(function, ptr_slice, slice_size, offset + slice_offset)); - - slice_offset += slice_size; - } + std::transform(slices.cbegin(), slices.cend(), std::back_inserter(slice_tasks), [&](auto& slice) { + return pool.submit(function, ptr + slice.offset, slice.size, offset + slice.offset); + }); return slice_tasks; } @@ -318,6 +311,21 @@ std::unique_ptr make_cufile_output(std::string const& filepa return nullptr; } +std::vector make_file_io_slices(size_t size, size_t max_slice_size) +{ + max_slice_size = std::max(1024ul, max_slice_size); + auto const n_slices = util::div_rounding_up_safe(size, max_slice_size); + std::vector slices; + slices.reserve(n_slices); + std::generate_n(std::back_inserter(slices), n_slices, [&, idx = 0]() mutable { + auto const slice_offset = idx++ * max_slice_size; + auto const slice_size = std::min(size - slice_offset, max_slice_size); + return file_io_slice{slice_offset, slice_size}; + }); + + return slices; +} + } // namespace detail } // namespace io } // namespace cudf diff --git a/cpp/src/io/utilities/file_io_utilities.hpp b/cpp/src/io/utilities/file_io_utilities.hpp index fcee4e43a20..be3ecc49ab0 100644 --- a/cpp/src/io/utilities/file_io_utilities.hpp +++ b/cpp/src/io/utilities/file_io_utilities.hpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021, NVIDIA CORPORATION. + * Copyright (c) 2021-2022, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -291,6 +291,21 @@ std::unique_ptr make_cufile_input(std::string const& filepath */ std::unique_ptr make_cufile_output(std::string const& filepath); +/** + * @brief Byte range to be read/written in a single operation. + */ +struct file_io_slice { + size_t offset; + size_t size; +}; + +/** + * @brief Split the total number of bytes to read/write into slices to enable parallel IO. + * + * If `max_slice_size` is below 1024, 1024 will be used instead to prevent potential misuse. + */ +std::vector make_file_io_slices(size_t size, size_t max_slice_size); + } // namespace detail } // namespace io } // namespace cudf diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 913761ecd03..27dd472b3f5 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -199,6 +199,7 @@ ConfigureTest( ConfigureTest(DECOMPRESSION_TEST io/comp/decomp_test.cpp) ConfigureTest(CSV_TEST io/csv_test.cpp) +ConfigureTest(FILE_IO_TEST io/file_io_test.cpp) ConfigureTest(ORC_TEST io/orc_test.cpp) ConfigureTest(PARQUET_TEST io/parquet_test.cpp) ConfigureTest(JSON_TEST io/json_test.cpp) diff --git a/cpp/tests/io/file_io_test.cpp b/cpp/tests/io/file_io_test.cpp new file mode 100644 index 00000000000..b546239fdca --- /dev/null +++ b/cpp/tests/io/file_io_test.cpp @@ -0,0 +1,46 @@ +/* + * Copyright (c) 2022, 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. + */ + +#include +#include + +#include + +#include + +// Base test fixture for tests +struct CuFileIOTest : public cudf::test::BaseFixture { +}; + +TEST_F(CuFileIOTest, SliceSize) +{ + std::vector> test_cases{ + {1 << 20, 1 << 18}, {1 << 18, 1 << 20}, {1 << 20, 3333}, {0, 1 << 18}, {0, 0}, {1 << 20, 0}}; + for (auto const& test_case : test_cases) { + auto const slices = cudf::io::detail::make_file_io_slices(test_case.first, test_case.second); + if (slices.empty()) { + ASSERT_EQ(test_case.first, 0); + } else { + ASSERT_EQ(slices.front().offset, 0); + ASSERT_EQ(slices.back().offset + slices.back().size, test_case.first); + for (auto i = 1u; i < slices.size(); ++i) { + ASSERT_EQ(slices[i].offset, slices[i - 1].offset + slices[i - 1].size); + } + } + } +} + +CUDF_TEST_PROGRAM_MAIN()