diff --git a/conda/recipes/libraft/build_libraft_distance.sh b/conda/recipes/libraft/build_libraft_distance.sh index 35bf354e9b..dca32b5238 100644 --- a/conda/recipes/libraft/build_libraft_distance.sh +++ b/conda/recipes/libraft/build_libraft_distance.sh @@ -1,4 +1,4 @@ #!/usr/bin/env bash -# Copyright (c) 2022, NVIDIA CORPORATION. +# Copyright (c) 2022-2023, NVIDIA CORPORATION. -./build.sh libraft -v --allgpuarch --compile-dist --no-nvtx +PARALLEL_LEVEL=8 ./build.sh libraft -v --allgpuarch --compile-dist --no-nvtx diff --git a/conda/recipes/libraft/build_libraft_nn.sh b/conda/recipes/libraft/build_libraft_nn.sh index 773d6ab02e..1d82e902a2 100644 --- a/conda/recipes/libraft/build_libraft_nn.sh +++ b/conda/recipes/libraft/build_libraft_nn.sh @@ -1,4 +1,4 @@ #!/usr/bin/env bash -# Copyright (c) 2022, NVIDIA CORPORATION. +# Copyright (c) 2022-2023, NVIDIA CORPORATION. -./build.sh libraft -v --allgpuarch --compile-nn --no-nvtx +PARALLEL_LEVEL=8 ./build.sh libraft -v --allgpuarch --compile-nn --no-nvtx diff --git a/conda/recipes/libraft/build_libraft_tests.sh b/conda/recipes/libraft/build_libraft_tests.sh index 040a2f8b8c..dc2fed2e6b 100644 --- a/conda/recipes/libraft/build_libraft_tests.sh +++ b/conda/recipes/libraft/build_libraft_tests.sh @@ -1,5 +1,5 @@ #!/usr/bin/env bash -# Copyright (c) 2022, NVIDIA CORPORATION. +# Copyright (c) 2022-2023, NVIDIA CORPORATION. -./build.sh tests bench -v --allgpuarch --no-nvtx +PARALLEL_LEVEL=8 ./build.sh tests bench -v --allgpuarch --no-nvtx cmake --install cpp/build --component testing diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index a6341f6dda..228e153f40 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -218,6 +218,15 @@ target_link_libraries( target_compile_features(raft INTERFACE cxx_std_17 $) +# Endian detection +include(TestBigEndian) +test_big_endian(BIG_ENDIAN) +if(BIG_ENDIAN) + target_compile_definitions(raft INTERFACE RAFT_SYSTEM_LITTLE_ENDIAN=0) +else() + target_compile_definitions(raft INTERFACE RAFT_SYSTEM_LITTLE_ENDIAN=1) +endif() + if(RAFT_COMPILE_DIST_LIBRARY OR RAFT_COMPILE_NN_LIBRARY) file( WRITE "${CMAKE_CURRENT_BINARY_DIR}/fatbin.ld" diff --git a/cpp/include/raft/core/detail/mdspan_numpy_serializer.hpp b/cpp/include/raft/core/detail/mdspan_numpy_serializer.hpp new file mode 100644 index 0000000000..df89811636 --- /dev/null +++ b/cpp/include/raft/core/detail/mdspan_numpy_serializer.hpp @@ -0,0 +1,487 @@ +/* + * Copyright (c) 2023, 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 +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace raft { + +namespace detail { + +namespace numpy_serializer { + +/** + * A small implementation of NumPy serialization format. + * Reference: https://numpy.org/doc/1.23/reference/generated/numpy.lib.format.html + * + * Adapted from https://github.com/llohse/libnpy/blob/master/include/npy.hpp, using the following + * license: + * + * MIT License + * + * Copyright (c) 2021 Leon Merten Lohse + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + */ + +#define RAFT_NUMPY_LITTLE_ENDIAN_CHAR '<' +#define RAFT_NUMPY_BIG_ENDIAN_CHAR '>' +#define RAFT_NUMPY_NO_ENDIAN_CHAR '|' +#define RAFT_NUMPY_MAGIC_STRING "\x93NUMPY" +#define RAFT_NUMPY_MAGIC_STRING_LENGTH 6 + +#if RAFT_SYSTEM_LITTLE_ENDIAN == 1 +#define RAFT_NUMPY_HOST_ENDIAN_CHAR RAFT_NUMPY_LITTLE_ENDIAN_CHAR +#else // RAFT_SYSTEM_LITTLE_ENDIAN == 1 +#define RAFT_NUMPY_HOST_ENDIAN_CHAR RAFT_NUMPY_BIG_ENDIAN_CHAR +#endif // RAFT_SYSTEM_LITTLE_ENDIAN == 1 + +using ndarray_len_t = std::uint64_t; + +struct dtype_t { + const char byteorder; + const char kind; + const unsigned int itemsize; + + std::string to_string() const + { + char buf[16] = {0}; + std::sprintf(buf, "%c%c%u", byteorder, kind, itemsize); + return std::string(buf); + } + + bool operator==(const dtype_t& other) const + { + return (byteorder == other.byteorder && kind == other.kind && itemsize == other.itemsize); + } +}; + +struct header_t { + const dtype_t dtype; + const bool fortran_order; + const std::vector shape; + + bool operator==(const header_t& other) const + { + return (dtype == other.dtype && fortran_order == other.fortran_order && shape == other.shape); + } +}; + +template +struct is_complex : std::false_type { +}; +template +struct is_complex> : std::true_type { +}; + +template , bool> = true> +inline dtype_t get_numpy_dtype() +{ + return {RAFT_NUMPY_HOST_ENDIAN_CHAR, 'f', sizeof(T)}; +} + +template && std::is_signed_v, bool> = true> +inline dtype_t get_numpy_dtype() +{ + const char endian_char = + (sizeof(T) == 1 ? RAFT_NUMPY_NO_ENDIAN_CHAR : RAFT_NUMPY_HOST_ENDIAN_CHAR); + return {endian_char, 'i', sizeof(T)}; +} + +template && std::is_unsigned_v, bool> = true> +inline dtype_t get_numpy_dtype() +{ + const char endian_char = + (sizeof(T) == 1 ? RAFT_NUMPY_NO_ENDIAN_CHAR : RAFT_NUMPY_HOST_ENDIAN_CHAR); + return {endian_char, 'u', sizeof(T)}; +} + +template {}, bool> = true> +inline dtype_t get_numpy_dtype() +{ + return {RAFT_NUMPY_HOST_ENDIAN_CHAR, 'c', sizeof(T)}; +} + +template , bool> = true> +inline dtype_t get_numpy_dtype() +{ + return get_numpy_dtype>(); +} + +template +inline std::string tuple_to_string(const std::vector& tuple) +{ + std::ostringstream oss; + if (tuple.empty()) { + oss << "()"; + } else if (tuple.size() == 1) { + oss << "(" << tuple.front() << ",)"; + } else { + oss << "("; + for (std::size_t i = 0; i < tuple.size() - 1; ++i) { + oss << tuple[i] << ", "; + } + oss << tuple.back() << ")"; + } + return oss.str(); +} + +inline std::string header_to_string(const header_t& header) +{ + std::ostringstream oss; + oss << "{'descr': '" << header.dtype.to_string() + << "', 'fortran_order': " << (header.fortran_order ? "True" : "False") + << ", 'shape': " << tuple_to_string(header.shape) << "}"; + return oss.str(); +} + +inline std::string trim(const std::string& str) +{ + const std::string whitespace = " \t"; + auto begin = str.find_first_not_of(whitespace); + if (begin == std::string::npos) { return ""; } + auto end = str.find_last_not_of(whitespace); + + return str.substr(begin, end - begin + 1); +} + +// A poor man's parser for Python dictionary +// TODO(hcho3): Consider writing a proper parser +// Limitation: can only parse a flat dictionary; all values are assumed to non-objects +// Limitation: must know all the keys ahead of time; you get undefined behavior if you omit any key +inline std::map parse_pydict(std::string str, + const std::vector& keys) +{ + std::map result; + + // Unwrap dictionary + str = trim(str); + RAFT_EXPECTS(str.front() == '{' && str.back() == '}', "Expected a Python dictionary"); + str = str.substr(1, str.length() - 2); + + // Get the position of each key and put it in the list + std::vector> positions; + for (auto const& key : keys) { + std::size_t pos = str.find("'" + key + "'"); + RAFT_EXPECTS(pos != std::string::npos, "Missing '%s' key.", key.c_str()); + positions.emplace_back(pos, key); + } + // Sort the list + std::sort(positions.begin(), positions.end()); + + // Extract each key-value pair + for (std::size_t i = 0; i < positions.size(); ++i) { + std::string key = positions[i].second; + + std::size_t begin = positions[i].first; + std::size_t end = (i + 1 < positions.size() ? positions[i + 1].first : std::string::npos); + std::string raw_value = trim(str.substr(begin, end - begin)); + if (raw_value.back() == ',') { raw_value.pop_back(); } + std::size_t sep_pos = raw_value.find_first_of(":"); + if (sep_pos == std::string::npos) { + result[key] = ""; + } else { + result[key] = trim(raw_value.substr(sep_pos + 1)); + } + } + + return result; +} + +inline std::string parse_pystring(std::string str) +{ + RAFT_EXPECTS(str.front() == '\'' && str.back() == '\'', "Invalid Python string: %s", str.c_str()); + return str.substr(1, str.length() - 2); +} + +inline bool parse_pybool(std::string str) +{ + if (str == "True") { + return true; + } else if (str == "False") { + return false; + } else { + RAFT_FAIL("Invalid Python boolean: %s", str.c_str()); + } +} + +inline std::vector parse_pytuple(std::string str) +{ + std::vector result; + + str = trim(str); + RAFT_EXPECTS(str.front() == '(' && str.back() == ')', "Invalid Python tuple: %s", str.c_str()); + str = str.substr(1, str.length() - 2); + + std::istringstream iss(str); + for (std::string token; std::getline(iss, token, ',');) { + result.push_back(trim(token)); + } + + return result; +} + +inline dtype_t parse_descr(std::string typestr) +{ + RAFT_EXPECTS(typestr.length() >= 3, "Invalid typestr: Too short"); + char byteorder_c = typestr.at(0); + char kind_c = typestr.at(1); + std::string itemsize_s = typestr.substr(2); + + const char endian_chars[] = { + RAFT_NUMPY_LITTLE_ENDIAN_CHAR, RAFT_NUMPY_BIG_ENDIAN_CHAR, RAFT_NUMPY_NO_ENDIAN_CHAR}; + const char numtype_chars[] = {'f', 'i', 'u', 'c'}; + + RAFT_EXPECTS(std::find(std::begin(endian_chars), std::end(endian_chars), byteorder_c) != + std::end(endian_chars), + "Invalid typestr: unrecognized byteorder %c", + byteorder_c); + RAFT_EXPECTS(std::find(std::begin(numtype_chars), std::end(numtype_chars), kind_c) != + std::end(numtype_chars), + "Invalid typestr: unrecognized kind %c", + kind_c); + unsigned int itemsize = std::stoul(itemsize_s); + + return {byteorder_c, kind_c, itemsize}; +} + +inline void write_magic(std::ostream& os) +{ + os.write(RAFT_NUMPY_MAGIC_STRING, RAFT_NUMPY_MAGIC_STRING_LENGTH); + RAFT_EXPECTS(os.good(), "Error writing magic string"); + // Use version 1.0 + os.put(1); + os.put(0); + RAFT_EXPECTS(os.good(), "Error writing magic string"); +} + +inline void read_magic(std::istream& is) +{ + char magic_buf[RAFT_NUMPY_MAGIC_STRING_LENGTH + 2] = {0}; + is.read(magic_buf, RAFT_NUMPY_MAGIC_STRING_LENGTH + 2); + RAFT_EXPECTS(is.good(), "Error reading magic string"); + + RAFT_EXPECTS(std::memcmp(magic_buf, RAFT_NUMPY_MAGIC_STRING, RAFT_NUMPY_MAGIC_STRING_LENGTH) == 0, + "The given stream does not have a valid NumPy format."); + + std::uint8_t version_major = magic_buf[RAFT_NUMPY_MAGIC_STRING_LENGTH]; + std::uint8_t version_minor = magic_buf[RAFT_NUMPY_MAGIC_STRING_LENGTH + 1]; + RAFT_EXPECTS(version_major == 1 && version_minor == 0, + "Unsupported NumPy version: %d.%d", + version_major, + version_minor); +} + +inline void write_header(std::ostream& os, const header_t& header) +{ + std::string header_dict = header_to_string(header); + std::size_t preamble_length = RAFT_NUMPY_MAGIC_STRING_LENGTH + 2 + 2 + header_dict.length() + 1; + RAFT_EXPECTS(preamble_length < 255 * 255, "Header too long"); + // Enforce 64-byte alignment + std::size_t padding_len = 64 - preamble_length % 64; + std::string padding(padding_len, ' '); + + write_magic(os); + + // Write header length + std::uint8_t header_len_le16[2]; + std::uint16_t header_len = + static_cast(header_dict.length() + padding.length() + 1); + header_len_le16[0] = (header_len >> 0) & 0xff; + header_len_le16[1] = (header_len >> 8) & 0xff; + os.put(header_len_le16[0]); + os.put(header_len_le16[1]); + RAFT_EXPECTS(os.good(), "Error writing HEADER_LEN"); + + os << header_dict << padding << "\n"; + RAFT_EXPECTS(os.good(), "Error writing header dict"); +} + +inline std::string read_header_bytes(std::istream& is) +{ + read_magic(is); + + // Read header length + std::uint8_t header_len_le16[2]; + is.read(reinterpret_cast(header_len_le16), 2); + RAFT_EXPECTS(is.good(), "Error while reading HEADER_LEN"); + const std::uint32_t header_length = (header_len_le16[0] << 0) | (header_len_le16[1] << 8); + + std::vector header_bytes(header_length); + is.read(header_bytes.data(), header_length); + RAFT_EXPECTS(is.good(), "Error while reading the header"); + + return std::string(header_bytes.data(), header_length); +} + +inline header_t read_header(std::istream& is) +{ + std::string header_bytes = read_header_bytes(is); + + // remove trailing newline + RAFT_EXPECTS(header_bytes.back() == '\n', "Invalid NumPy header"); + header_bytes.pop_back(); + + // parse the header dict + auto header_dict = parse_pydict(header_bytes, {"descr", "fortran_order", "shape"}); + dtype_t descr = parse_descr(parse_pystring(header_dict["descr"])); + bool fortran_order = parse_pybool(header_dict["fortran_order"]); + std::vector shape; + auto shape_tup_str = parse_pytuple(header_dict["shape"]); + for (const auto& e : shape_tup_str) { + shape.push_back(static_cast(std::stoul(e))); + } + + RAFT_EXPECTS( + descr.byteorder == RAFT_NUMPY_HOST_ENDIAN_CHAR || descr.byteorder == RAFT_NUMPY_NO_ENDIAN_CHAR, + "The mdspan was serialized on a %s machine but you're attempting to load it on " + "a %s machine. This use case is not currently supported.", + (RAFT_SYSTEM_LITTLE_ENDIAN ? "big-endian" : "little-endian"), + (RAFT_SYSTEM_LITTLE_ENDIAN ? "little-endian" : "big-endian")); + + return {descr, fortran_order, shape}; +} + +template +inline void serialize_host_mdspan( + std::ostream& os, + const raft::host_mdspan& obj) +{ + static_assert(std::is_same_v || + std::is_same_v, + "The serializer only supports row-major and column-major layouts"); + + using obj_t = raft::host_mdspan; + + const auto dtype = get_numpy_dtype(); + const bool fortran_order = std::is_same_v; + std::vector shape; + for (typename obj_t::rank_type i = 0; i < obj.rank(); ++i) { + shape.push_back(obj.extent(i)); + } + const header_t header = {dtype, fortran_order, shape}; + write_header(os, header); + + // For contiguous layouts, size() == product of dimensions + os.write(reinterpret_cast(obj.data_handle()), obj.size() * sizeof(ElementType)); + RAFT_EXPECTS(os.good(), "Error writing content of mdspan"); +} + +template +inline void serialize_scalar(std::ostream& os, const T& value) +{ + const auto dtype = get_numpy_dtype(); + const bool fortran_order = false; + const std::vector shape{}; + const header_t header = {dtype, fortran_order, shape}; + write_header(os, header); + os.write(reinterpret_cast(&value), sizeof(T)); + RAFT_EXPECTS(os.good(), "Error serializing a scalar"); +} + +template +inline void deserialize_host_mdspan( + std::istream& is, + const raft::host_mdspan& obj) +{ + static_assert(std::is_same_v || + std::is_same_v, + "The serializer only supports row-major and column-major layouts"); + + using obj_t = raft::host_mdspan; + + // Check if given dtype and fortran_order are compatible with the mdspan + const auto expected_dtype = get_numpy_dtype(); + const bool expected_fortran_order = std::is_same_v; + header_t header = read_header(is); + RAFT_EXPECTS(header.dtype == expected_dtype, + "Expected dtype %s but got %s instead", + header.dtype.to_string().c_str(), + expected_dtype.to_string().c_str()); + RAFT_EXPECTS(header.fortran_order == expected_fortran_order, + "Wrong matrix layout; expected %s but got a different layout", + (expected_fortran_order ? "Fortran layout" : "C layout")); + + // Check if dimensions are correct + RAFT_EXPECTS(obj.rank() == header.shape.size(), + "Incorrect rank: expected %zu but got %zu", + obj.rank(), + header.shape.size()); + for (typename obj_t::rank_type i = 0; i < obj.rank(); ++i) { + RAFT_EXPECTS(static_cast(obj.extent(i)) == header.shape[i], + "Incorrect dimension: expected %zu but got %zu", + static_cast(obj.extent(i)), + header.shape[i]); + } + + // For contiguous layouts, size() == product of dimensions + is.read(reinterpret_cast(obj.data_handle()), obj.size() * sizeof(ElementType)); + RAFT_EXPECTS(is.good(), "Error while reading mdspan content"); +} + +template +inline T deserialize_scalar(std::istream& is) +{ + // Check if dtype is correct + const auto expected_dtype = get_numpy_dtype(); + header_t header = read_header(is); + RAFT_EXPECTS(header.dtype == expected_dtype, + "Expected dtype %s but got %s instead", + header.dtype.to_string().c_str(), + expected_dtype.to_string().c_str()); + // Check if dimensions are correct; shape should be () + RAFT_EXPECTS(header.shape.empty(), "Incorrect rank: expected 0 but got %zu", header.shape.size()); + + T value; + is.read(reinterpret_cast(&value), sizeof(T)); + RAFT_EXPECTS(is.good(), "Error while deserializing scalar"); + return value; +} + +} // end namespace numpy_serializer +} // end namespace detail +} // end namespace raft diff --git a/cpp/include/raft/core/serialize.hpp b/cpp/include/raft/core/serialize.hpp new file mode 100644 index 0000000000..05814e2845 --- /dev/null +++ b/cpp/include/raft/core/serialize.hpp @@ -0,0 +1,167 @@ +/* + * Copyright (c) 2023, 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 +#include +#include +#include + +#include +#include + +/** + * Collection of serialization functions for RAFT data types + */ + +namespace raft { + +template +inline void serialize_mdspan( + const raft::device_resources&, + std::ostream& os, + const raft::host_mdspan& obj) +{ + detail::numpy_serializer::serialize_host_mdspan(os, obj); +} + +template +inline void serialize_mdspan( + const raft::device_resources& handle, + std::ostream& os, + const raft::device_mdspan& obj) +{ + static_assert(std::is_same_v || + std::is_same_v, + "The serializer only supports row-major and column-major layouts"); + using obj_t = raft::device_mdspan; + + // Copy to host before serializing + // For contiguous layouts, size() == product of dimensions + std::vector tmp(obj.size()); + cudaStream_t stream = handle.get_stream(); + raft::update_host(tmp.data(), obj.data_handle(), obj.size(), stream); + handle.sync_stream(); + using inner_accessor_type = typename obj_t::accessor_type::accessor_type; + auto tmp_mdspan = + raft::host_mdspan>( + tmp.data(), obj.extents()); + detail::numpy_serializer::serialize_host_mdspan(os, tmp_mdspan); +} + +template +inline void serialize_mdspan( + const raft::device_resources&, + std::ostream& os, + const raft::managed_mdspan& obj) +{ + using obj_t = raft::managed_mdspan; + using inner_accessor_type = typename obj_t::accessor_type::accessor_type; + auto tmp_mdspan = + raft::host_mdspan>( + obj.data_handle(), obj.extents()); + detail::numpy_serializer::serialize_host_mdspan(os, tmp_mdspan); +} + +template +inline void deserialize_mdspan( + const raft::device_resources&, + std::istream& is, + raft::host_mdspan& obj) +{ + detail::numpy_serializer::deserialize_host_mdspan(is, obj); +} + +template +inline void deserialize_mdspan( + const raft::device_resources& handle, + std::istream& is, + raft::device_mdspan& obj) +{ + static_assert(std::is_same_v || + std::is_same_v, + "The serializer only supports row-major and column-major layouts"); + using obj_t = raft::device_mdspan; + + // Copy to device after serializing + // For contiguous layouts, size() == product of dimensions + std::vector tmp(obj.size()); + using inner_accessor_type = typename obj_t::accessor_type::accessor_type; + auto tmp_mdspan = + raft::host_mdspan>( + tmp.data(), obj.extents()); + detail::numpy_serializer::deserialize_host_mdspan(is, tmp_mdspan); + + cudaStream_t stream = handle.get_stream(); + raft::update_device(obj.data_handle(), tmp.data(), obj.size(), stream); + handle.sync_stream(); +} + +template +inline void deserialize_mdspan( + const raft::device_resources& handle, + std::istream& is, + raft::host_mdspan&& obj) +{ + deserialize_mdspan(handle, is, obj); +} + +template +inline void deserialize_mdspan( + const raft::device_resources& handle, + std::istream& is, + raft::managed_mdspan& obj) +{ + using obj_t = raft::managed_mdspan; + using inner_accessor_type = typename obj_t::accessor_type::accessor_type; + auto tmp_mdspan = + raft::host_mdspan>( + obj.data_handle(), obj.extents()); + detail::numpy_serializer::deserialize_host_mdspan(is, tmp_mdspan); +} + +template +inline void deserialize_mdspan( + const raft::device_resources& handle, + std::istream& is, + raft::managed_mdspan&& obj) +{ + deserialize_mdspan(handle, is, obj); +} + +template +inline void deserialize_mdspan( + const raft::device_resources& handle, + std::istream& is, + raft::device_mdspan&& obj) +{ + deserialize_mdspan(handle, is, obj); +} + +template +inline void serialize_scalar(const raft::device_resources&, std::ostream& os, const T& value) +{ + detail::numpy_serializer::serialize_scalar(os, value); +} + +template +inline T deserialize_scalar(const raft::device_resources&, std::istream& is) +{ + return detail::numpy_serializer::deserialize_scalar(is); +} + +} // end namespace raft diff --git a/cpp/include/raft/spatial/knn/detail/ann_serialization.h b/cpp/include/raft/spatial/knn/detail/ann_serialization.h deleted file mode 100644 index 34dc3cf00f..0000000000 --- a/cpp/include/raft/spatial/knn/detail/ann_serialization.h +++ /dev/null @@ -1,140 +0,0 @@ -/* - * Copyright (c) 2022-2023, 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 -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace raft::spatial::knn::detail { - -template -void write_scalar(std::ofstream& of, const T& value) -{ - of.write((char*)&value, sizeof value); - if (of.good()) { - RAFT_LOG_DEBUG("Written %z bytes", (sizeof value)); - } else { - RAFT_FAIL("error writing value to file"); - } -} - -template -T read_scalar(std::ifstream& file) -{ - T value; - file.read((char*)&value, sizeof value); - if (file.good()) { - RAFT_LOG_DEBUG("Read %z bytes", (sizeof value)); - } else { - RAFT_FAIL("error reading value from file"); - } - return value; -} - -template -void write_mdspan( - raft::device_resources const& handle, - std::ofstream& of, - const raft::device_mdspan& obj) -{ - using obj_t = raft::device_mdspan; - write_scalar(of, obj.rank()); - if (obj.is_exhaustive() && obj.is_unique()) { - write_scalar(of, obj.size()); - } else { - RAFT_FAIL("Cannot serialize non exhaustive mdarray"); - } - if (obj.size() > 0) { - for (typename obj_t::rank_type i = 0; i < obj.rank(); i++) - write_scalar(of, obj.extent(i)); - cudaStream_t stream = handle.get_stream(); - std::vector< - typename raft::device_mdspan::value_type> - tmp(obj.size()); - raft::update_host(tmp.data(), obj.data_handle(), obj.size(), stream); - handle.sync_stream(stream); - of.write(reinterpret_cast(tmp.data()), tmp.size() * sizeof(ElementType)); - if (of.good()) { - RAFT_LOG_DEBUG("Written %zu bytes", - static_cast(obj.size() * sizeof(obj.data_handle()[0]))); - } else { - RAFT_FAIL("Error writing mdarray to file"); - } - } else { - RAFT_LOG_DEBUG("Skipping mdspand with zero size"); - } -} - -template -void read_mdspan(raft::device_resources const& handle, - std::ifstream& file, - raft::device_mdspan& obj) -{ - using obj_t = raft::device_mdspan; - auto rank = read_scalar(file); - if (obj.rank() != rank) { RAFT_FAIL("Incorrect rank while reading mdarray"); } - auto size = read_scalar(file); - if (obj.size() != size) { - RAFT_FAIL("Incorrect rank while reading mdarray %zu vs %zu", - static_cast(size), - static_cast(obj.size())); - } - if (obj.size() > 0) { - for (typename obj_t::rank_type i = 0; i < obj.rank(); i++) { - auto ex = read_scalar(file); - if (obj.extent(i) != ex) { - RAFT_FAIL("Incorrect extent while reading mdarray %d vs %d at %d", - static_cast(ex), - static_cast(obj.extent(i)), - static_cast(i)); - } - } - cudaStream_t stream = handle.get_stream(); - std::vector tmp(obj.size()); - file.read(reinterpret_cast(tmp.data()), tmp.size() * sizeof(ElementType)); - raft::update_device(obj.data_handle(), tmp.data(), tmp.size(), stream); - handle.sync_stream(stream); - if (file.good()) { - RAFT_LOG_DEBUG("Read %zu bytes", - static_cast(obj.size() * sizeof(obj.data_handle()[0]))); - } else { - RAFT_FAIL("error reading mdarray from file"); - } - } else { - RAFT_LOG_DEBUG("Skipping mdspand with zero size"); - } -} - -template -void read_mdspan(raft::device_resources const& handle, - std::ifstream& file, - raft::device_mdspan&& obj) -{ - read_mdspan(handle, file, obj); -} -} // namespace raft::spatial::knn::detail diff --git a/cpp/include/raft/spatial/knn/detail/ivf_flat_build.cuh b/cpp/include/raft/spatial/knn/detail/ivf_flat_build.cuh index 08a7a461a4..c417a97531 100644 --- a/cpp/include/raft/spatial/knn/detail/ivf_flat_build.cuh +++ b/cpp/include/raft/spatial/knn/detail/ivf_flat_build.cuh @@ -17,7 +17,6 @@ #pragma once #include "../ivf_flat_types.hpp" -#include "ann_serialization.h" #include "ann_utils.cuh" #include @@ -26,6 +25,7 @@ #include #include #include +#include #include #include #include @@ -34,6 +34,9 @@ #include +#include +#include + namespace raft::spatial::knn::ivf_flat::detail { using namespace raft::spatial::knn::detail; // NOLINT @@ -386,7 +389,16 @@ inline void fill_refinement_index(raft::device_resources const& handle, RAFT_CUDA_TRY(cudaPeekAtLastError()); } -static const int serialization_version = 1; +// Serialization version 2 +// No backward compatibility yet; that is, can't add additional fields without breaking +// backward compatibility. +// TODO(hcho3) Implement next-gen serializer for IVF that allows for expansion in a backward +// compatible fashion. +constexpr int serialization_version = 2; + +static_assert(sizeof(index) == 408, + "The size of the index struct has changed since the last update; " + "paste in the new size and consider updating the save/load logic"); /** * Save the index to file. @@ -399,34 +411,34 @@ static const int serialization_version = 1; * */ template -void save(raft::device_resources const& handle, - const std::string& filename, - const index& index_) +void serialize(raft::device_resources const& handle, + const std::string& filename, + const index& index_) { std::ofstream of(filename, std::ios::out | std::ios::binary); if (!of) { RAFT_FAIL("Cannot open %s", filename.c_str()); } RAFT_LOG_DEBUG( "Saving IVF-PQ index, size %zu, dim %u", static_cast(index_.size()), index_.dim()); - write_scalar(of, serialization_version); - write_scalar(of, index_.size()); - write_scalar(of, index_.dim()); - write_scalar(of, index_.n_lists()); - write_scalar(of, index_.metric()); - write_scalar(of, index_.veclen()); - write_scalar(of, index_.adaptive_centers()); - write_mdspan(handle, of, index_.data()); - write_mdspan(handle, of, index_.indices()); - write_mdspan(handle, of, index_.list_sizes()); - write_mdspan(handle, of, index_.list_offsets()); - write_mdspan(handle, of, index_.centers()); + serialize_scalar(handle, of, serialization_version); + serialize_scalar(handle, of, index_.size()); + serialize_scalar(handle, of, index_.dim()); + serialize_scalar(handle, of, index_.n_lists()); + serialize_scalar(handle, of, index_.metric()); + serialize_scalar(handle, of, index_.veclen()); + serialize_scalar(handle, of, index_.adaptive_centers()); + serialize_mdspan(handle, of, index_.data()); + serialize_mdspan(handle, of, index_.indices()); + serialize_mdspan(handle, of, index_.list_sizes()); + serialize_mdspan(handle, of, index_.list_offsets()); + serialize_mdspan(handle, of, index_.centers()); if (index_.center_norms()) { bool has_norms = true; - write_scalar(of, has_norms); - write_mdspan(handle, of, *index_.center_norms()); + serialize_scalar(handle, of, has_norms); + serialize_mdspan(handle, of, *index_.center_norms()); } else { bool has_norms = false; - write_scalar(of, has_norms); + serialize_scalar(handle, of, has_norms); } of.close(); if (!of) { RAFT_FAIL("Error writing output %s", filename.c_str()); } @@ -442,40 +454,41 @@ void save(raft::device_resources const& handle, * */ template -auto load(raft::device_resources const& handle, const std::string& filename) -> index +auto deserialize(raft::device_resources const& handle, const std::string& filename) + -> index { std::ifstream infile(filename, std::ios::in | std::ios::binary); if (!infile) { RAFT_FAIL("Cannot open %s", filename.c_str()); } - auto ver = read_scalar(infile); + auto ver = deserialize_scalar(handle, infile); if (ver != serialization_version) { RAFT_FAIL("serialization version mismatch, expected %d, got %d ", serialization_version, ver); } - auto n_rows = read_scalar(infile); - auto dim = read_scalar(infile); - auto n_lists = read_scalar(infile); - auto metric = read_scalar(infile); - auto veclen = read_scalar(infile); - bool adaptive_centers = read_scalar(infile); + auto n_rows = deserialize_scalar(handle, infile); + auto dim = deserialize_scalar(handle, infile); + auto n_lists = deserialize_scalar(handle, infile); + auto metric = deserialize_scalar(handle, infile); + auto veclen = deserialize_scalar(handle, infile); + bool adaptive_centers = deserialize_scalar(handle, infile); index index_ = raft::spatial::knn::ivf_flat::index(handle, metric, n_lists, adaptive_centers, dim); index_.allocate(handle, n_rows); auto data = index_.data(); - read_mdspan(handle, infile, data); - read_mdspan(handle, infile, index_.indices()); - read_mdspan(handle, infile, index_.list_sizes()); - read_mdspan(handle, infile, index_.list_offsets()); - read_mdspan(handle, infile, index_.centers()); - bool has_norms = read_scalar(infile); + deserialize_mdspan(handle, infile, data); + deserialize_mdspan(handle, infile, index_.indices()); + deserialize_mdspan(handle, infile, index_.list_sizes()); + deserialize_mdspan(handle, infile, index_.list_offsets()); + deserialize_mdspan(handle, infile, index_.centers()); + bool has_norms = deserialize_scalar(handle, infile); if (has_norms) { if (!index_.center_norms()) { RAFT_FAIL("Error inconsistent center norms"); } else { auto center_norms = *index_.center_norms(); - read_mdspan(handle, infile, center_norms); + deserialize_mdspan(handle, infile, center_norms); } } infile.close(); diff --git a/cpp/include/raft/spatial/knn/detail/ivf_pq_build.cuh b/cpp/include/raft/spatial/knn/detail/ivf_pq_build.cuh index adc485d3bf..66a4207b20 100644 --- a/cpp/include/raft/spatial/knn/detail/ivf_pq_build.cuh +++ b/cpp/include/raft/spatial/knn/detail/ivf_pq_build.cuh @@ -16,7 +16,6 @@ #pragma once -#include "ann_serialization.h" #include "ann_utils.cuh" #include @@ -27,6 +26,7 @@ #include #include #include +#include #include #include #include @@ -55,6 +55,8 @@ #include #include +#include +#include #include namespace raft::spatial::knn::ivf_pq::detail { @@ -1388,7 +1390,16 @@ auto build(raft::device_resources const& handle, } } -static const int serialization_version = 1; +// Serialization version 2 +// No backward compatibility yet; that is, can't add additional fields without breaking +// backward compatibility. +// TODO(hcho3) Implement next-gen serializer for IVF that allows for expansion in a backward +// compatible fashion. +constexpr int serialization_version = 2; + +static_assert(sizeof(index) == 560, + "The size of the index struct has changed since the last update; " + "paste in the new size and consider updating the save/load logic"); /** * Save the index to file. @@ -1401,9 +1412,9 @@ static const int serialization_version = 1; * */ template -void save(raft::device_resources const& handle_, - const std::string& filename, - const index& index_) +void serialize(raft::device_resources const& handle_, + const std::string& filename, + const index& index_) { std::ofstream of(filename, std::ios::out | std::ios::binary); if (!of) { RAFT_FAIL("Cannot open file %s", filename.c_str()); } @@ -1414,25 +1425,25 @@ void save(raft::device_resources const& handle_, static_cast(index_.pq_dim()), static_cast(index_.pq_bits())); - write_scalar(of, serialization_version); - write_scalar(of, index_.size()); - write_scalar(of, index_.dim()); - write_scalar(of, index_.pq_bits()); - write_scalar(of, index_.pq_dim()); - - write_scalar(of, index_.metric()); - write_scalar(of, index_.codebook_kind()); - write_scalar(of, index_.n_lists()); - write_scalar(of, index_.n_nonempty_lists()); - - write_mdspan(handle_, of, index_.pq_centers()); - write_mdspan(handle_, of, index_.pq_dataset()); - write_mdspan(handle_, of, index_.indices()); - write_mdspan(handle_, of, index_.rotation_matrix()); - write_mdspan(handle_, of, index_.list_offsets()); - write_mdspan(handle_, of, index_.list_sizes()); - write_mdspan(handle_, of, index_.centers()); - write_mdspan(handle_, of, index_.centers_rot()); + serialize_scalar(handle_, of, serialization_version); + serialize_scalar(handle_, of, index_.size()); + serialize_scalar(handle_, of, index_.dim()); + serialize_scalar(handle_, of, index_.pq_bits()); + serialize_scalar(handle_, of, index_.pq_dim()); + + serialize_scalar(handle_, of, index_.metric()); + serialize_scalar(handle_, of, index_.codebook_kind()); + serialize_scalar(handle_, of, index_.n_lists()); + serialize_scalar(handle_, of, index_.n_nonempty_lists()); + + serialize_mdspan(handle_, of, index_.pq_centers()); + serialize_mdspan(handle_, of, index_.pq_dataset()); + serialize_mdspan(handle_, of, index_.indices()); + serialize_mdspan(handle_, of, index_.rotation_matrix()); + serialize_mdspan(handle_, of, index_.list_offsets()); + serialize_mdspan(handle_, of, index_.list_sizes()); + serialize_mdspan(handle_, of, index_.centers()); + serialize_mdspan(handle_, of, index_.centers_rot()); of.close(); if (!of) { RAFT_FAIL("Error writing output %s", filename.c_str()); } @@ -1450,28 +1461,28 @@ void save(raft::device_resources const& handle_, * */ template -auto load(raft::device_resources const& handle_, const std::string& filename) -> index +auto deserialize(raft::device_resources const& handle_, const std::string& filename) -> index { std::ifstream infile(filename, std::ios::in | std::ios::binary); if (!infile) { RAFT_FAIL("Cannot open file %s", filename.c_str()); } - auto ver = read_scalar(infile); + auto ver = deserialize_scalar(handle_, infile); if (ver != serialization_version) { RAFT_FAIL("serialization version mismatch %d vs. %d", ver, serialization_version); } - auto n_rows = read_scalar(infile); - auto dim = read_scalar(infile); - auto pq_bits = read_scalar(infile); - auto pq_dim = read_scalar(infile); + auto n_rows = deserialize_scalar(handle_, infile); + auto dim = deserialize_scalar(handle_, infile); + auto pq_bits = deserialize_scalar(handle_, infile); + auto pq_dim = deserialize_scalar(handle_, infile); - auto metric = read_scalar(infile); - auto codebook_kind = read_scalar(infile); - auto n_lists = read_scalar(infile); - auto n_nonempty_lists = read_scalar(infile); + auto metric = deserialize_scalar(handle_, infile); + auto codebook_kind = deserialize_scalar(handle_, infile); + auto n_lists = deserialize_scalar(handle_, infile); + auto n_nonempty_lists = deserialize_scalar(handle_, infile); RAFT_LOG_DEBUG("n_rows %zu, dim %d, pq_dim %d, pq_bits %d, n_lists %d", - static_cast(n_rows), + static_cast(n_rows), static_cast(dim), static_cast(pq_dim), static_cast(pq_bits), @@ -1481,14 +1492,14 @@ auto load(raft::device_resources const& handle_, const std::string& filename) -> handle_, metric, codebook_kind, n_lists, dim, pq_bits, pq_dim, n_nonempty_lists); index_.allocate(handle_, n_rows); - read_mdspan(handle_, infile, index_.pq_centers()); - read_mdspan(handle_, infile, index_.pq_dataset()); - read_mdspan(handle_, infile, index_.indices()); - read_mdspan(handle_, infile, index_.rotation_matrix()); - read_mdspan(handle_, infile, index_.list_offsets()); - read_mdspan(handle_, infile, index_.list_sizes()); - read_mdspan(handle_, infile, index_.centers()); - read_mdspan(handle_, infile, index_.centers_rot()); + deserialize_mdspan(handle_, infile, index_.pq_centers()); + deserialize_mdspan(handle_, infile, index_.pq_dataset()); + deserialize_mdspan(handle_, infile, index_.indices()); + deserialize_mdspan(handle_, infile, index_.rotation_matrix()); + deserialize_mdspan(handle_, infile, index_.list_offsets()); + deserialize_mdspan(handle_, infile, index_.list_sizes()); + deserialize_mdspan(handle_, infile, index_.centers()); + deserialize_mdspan(handle_, infile, index_.centers_rot()); infile.close(); diff --git a/cpp/include/raft/util/cudart_utils.hpp b/cpp/include/raft/util/cudart_utils.hpp index 7090d4d2bf..1c9793eb0a 100644 --- a/cpp/include/raft/util/cudart_utils.hpp +++ b/cpp/include/raft/util/cudart_utils.hpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2022, NVIDIA CORPORATION. + * Copyright (c) 2019-2023, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. diff --git a/cpp/include/raft/util/device_atomics.cuh b/cpp/include/raft/util/device_atomics.cuh index 6e956e8e38..14856bed8e 100644 --- a/cpp/include/raft/util/device_atomics.cuh +++ b/cpp/include/raft/util/device_atomics.cuh @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2022, NVIDIA CORPORATION. + * Copyright (c) 2019-2023, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. diff --git a/cpp/include/raft_runtime/neighbors/ivf_pq.hpp b/cpp/include/raft_runtime/neighbors/ivf_pq.hpp index cae32c9530..59d0b59128 100644 --- a/cpp/include/raft_runtime/neighbors/ivf_pq.hpp +++ b/cpp/include/raft_runtime/neighbors/ivf_pq.hpp @@ -84,9 +84,9 @@ RAFT_INST_BUILD_EXTEND(uint8_t, uint64_t) * @param[in] index IVF-PQ index * */ -void save(raft::device_resources const& handle, - const std::string& filename, - const raft::neighbors::ivf_pq::index& index); +void serialize(raft::device_resources const& handle, + const std::string& filename, + const raft::neighbors::ivf_pq::index& index); /** * Load index from file. @@ -98,8 +98,8 @@ void save(raft::device_resources const& handle, * @param[in] index IVF-PQ index * */ -void load(raft::device_resources const& handle, - const std::string& filename, - raft::neighbors::ivf_pq::index* index); +void deserialize(raft::device_resources const& handle, + const std::string& filename, + raft::neighbors::ivf_pq::index* index); } // namespace raft::runtime::neighbors::ivf_pq diff --git a/cpp/src/distance/neighbors/ivfpq_build.cu b/cpp/src/distance/neighbors/ivfpq_build.cu index 650767f918..31e304835b 100644 --- a/cpp/src/distance/neighbors/ivfpq_build.cu +++ b/cpp/src/distance/neighbors/ivfpq_build.cu @@ -64,18 +64,18 @@ RAFT_INST_BUILD_EXTEND(uint8_t, uint64_t); #undef RAFT_INST_BUILD_EXTEND -void save(raft::device_resources const& handle, - const std::string& filename, - const raft::neighbors::ivf_pq::index& index) +void serialize(raft::device_resources const& handle, + const std::string& filename, + const raft::neighbors::ivf_pq::index& index) { - raft::spatial::knn::ivf_pq::detail::save(handle, filename, index); + raft::spatial::knn::ivf_pq::detail::serialize(handle, filename, index); }; -void load(raft::device_resources const& handle, - const std::string& filename, - raft::neighbors::ivf_pq::index* index) +void deserialize(raft::device_resources const& handle, + const std::string& filename, + raft::neighbors::ivf_pq::index* index) { if (!index) { RAFT_FAIL("Invalid index pointer"); } - *index = raft::spatial::knn::ivf_pq::detail::load(handle, filename); + *index = raft::spatial::knn::ivf_pq::detail::deserialize(handle, filename); }; } // namespace raft::runtime::neighbors::ivf_pq diff --git a/cpp/test/CMakeLists.txt b/cpp/test/CMakeLists.txt index af8bf844df..2e89418f8e 100644 --- a/cpp/test/CMakeLists.txt +++ b/cpp/test/CMakeLists.txt @@ -96,6 +96,7 @@ if(BUILD_TESTS) test/core/nvtx.cpp test/core/mdarray.cu test/core/mdspan_utils.cu + test/core/numpy_serializer.cu test/core/memory_type.cpp test/core/span.cpp test/core/span.cu diff --git a/cpp/test/core/mdspan_utils.cu b/cpp/test/core/mdspan_utils.cu index 4bb689c8c0..448391fa95 100644 --- a/cpp/test/core/mdspan_utils.cu +++ b/cpp/test/core/mdspan_utils.cu @@ -241,4 +241,4 @@ void test_const_mdspan() TEST(MDSpan, ConstMDSpan) { test_const_mdspan(); } -} // namespace raft \ No newline at end of file +} // namespace raft diff --git a/cpp/test/core/numpy_serializer.cu b/cpp/test/core/numpy_serializer.cu new file mode 100644 index 0000000000..4131a33171 --- /dev/null +++ b/cpp/test/core/numpy_serializer.cu @@ -0,0 +1,271 @@ +/* + * Copyright (c) 2023, 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 + +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace { + +template +using dextents = std::experimental::dextents; + +} // anonymous namespace + +namespace raft { + +template +void test_mdspan_roundtrip(const raft::device_resources& handle, VectorType& vec, Args... dims) +{ + VectorType vec2(vec.size()); + + auto span = MDSpanType(thrust::raw_pointer_cast(vec.data()), dims...); + std::ostringstream oss; + serialize_mdspan(handle, oss, span); + + auto span2 = MDSpanType(thrust::raw_pointer_cast(vec2.data()), dims...); + std::istringstream iss(oss.str()); + deserialize_mdspan(handle, iss, span2); + EXPECT_EQ(vec, vec2); +} + +template +void run_roundtrip_test_mdspan_serializer() +{ + raft::device_resources handle{}; + thrust::host_vector vec = std::vector{1, 2, 3, 4, 5, 6, 7, 8}; + + using mdspan_matrix2d_c_layout = + raft::host_mdspan, raft::layout_c_contiguous>; + using mdspan_matrix2d_f_layout = + raft::host_mdspan, raft::layout_f_contiguous>; + + test_mdspan_roundtrip(handle, vec, 2, 4); + test_mdspan_roundtrip(handle, vec, 2, 4); + + using device_mdspan_matrix3d_c_layout = + raft::device_mdspan, raft::layout_c_contiguous>; + using device_mdspan_matrix3d_f_layout = + raft::device_mdspan, raft::layout_f_contiguous>; + + thrust::device_vector d_vec(vec); + test_mdspan_roundtrip(handle, d_vec, 2, 2, 2); + test_mdspan_roundtrip(handle, d_vec, 2, 2, 2); +} + +TEST(NumPySerializerMDSpan, E2ERoundTrip) +{ + run_roundtrip_test_mdspan_serializer(); + run_roundtrip_test_mdspan_serializer(); + run_roundtrip_test_mdspan_serializer(); + run_roundtrip_test_mdspan_serializer(); + run_roundtrip_test_mdspan_serializer>(); +} + +TEST(NumPySerializerMDSpan, HeaderRoundTrip) +{ + char byteorder = RAFT_NUMPY_HOST_ENDIAN_CHAR; + for (char kind : std::vector{'f', 'i', 'u', 'c'}) { + for (unsigned int itemsize : std::vector{1, 2, 4, 8, 16}) { + for (bool fortran_order : std::vector{true, false}) { + for (const auto& shape : std::vector>{ + {10}, {2, 2}, {10, 30, 100}, {}}) { + detail::numpy_serializer::dtype_t dtype{byteorder, kind, itemsize}; + detail::numpy_serializer::header_t header{dtype, fortran_order, shape}; + std::ostringstream oss; + detail::numpy_serializer::write_header(oss, header); + std::istringstream iss(oss.str()); + auto header2 = detail::numpy_serializer::read_header(iss); + EXPECT_EQ(header, header2); + } + } + } + } +} + +TEST(NumPySerializerMDSpan, ManagedMDSpan) +{ + raft::device_resources handle{}; + thrust::universal_vector vec = std::vector{1, 2, 3, 4, 5, 6, 7, 8}; + using managed_mdspan_matrix2d_c_layout = + raft::managed_mdspan, raft::layout_c_contiguous>; + test_mdspan_roundtrip(handle, vec, 2, 2, 2); +} + +TEST(NumPySerializerMDSpan, Tuple2String) +{ + { + std::vector tuple{}; + EXPECT_EQ(detail::numpy_serializer::tuple_to_string(tuple), "()"); + } + { + std::vector tuple{2}; + EXPECT_EQ(detail::numpy_serializer::tuple_to_string(tuple), "(2,)"); + } + { + std::vector tuple{2, 3}; + EXPECT_EQ(detail::numpy_serializer::tuple_to_string(tuple), "(2, 3)"); + } + { + std::vector tuple{2, 3, 10, 20}; + EXPECT_EQ(detail::numpy_serializer::tuple_to_string(tuple), "(2, 3, 10, 20)"); + } +} + +TEST(NumPySerializerMDSpan, NumPyDType) +{ + const char expected_endian_char = RAFT_SYSTEM_LITTLE_ENDIAN ? '<' : '>'; + { + const detail::numpy_serializer::dtype_t expected_dtype{ + expected_endian_char, 'f', sizeof(float)}; + EXPECT_EQ(detail::numpy_serializer::get_numpy_dtype(), expected_dtype); + } + { + const detail::numpy_serializer::dtype_t expected_dtype{ + expected_endian_char, 'f', sizeof(long double)}; + EXPECT_EQ(detail::numpy_serializer::get_numpy_dtype(), expected_dtype); + } + { + const detail::numpy_serializer::dtype_t expected_dtype{'|', 'i', sizeof(signed char)}; + EXPECT_EQ(detail::numpy_serializer::get_numpy_dtype(), expected_dtype); + } + { + const detail::numpy_serializer::dtype_t expected_dtype{ + expected_endian_char, 'i', sizeof(std::int64_t)}; + EXPECT_EQ(detail::numpy_serializer::get_numpy_dtype(), expected_dtype); + } + { + const detail::numpy_serializer::dtype_t expected_dtype{'|', 'u', sizeof(unsigned char)}; + EXPECT_EQ(detail::numpy_serializer::get_numpy_dtype(), expected_dtype); + } + { + const detail::numpy_serializer::dtype_t expected_dtype{ + expected_endian_char, 'u', sizeof(std::uint64_t)}; + EXPECT_EQ(detail::numpy_serializer::get_numpy_dtype(), expected_dtype); + } + { + const detail::numpy_serializer::dtype_t expected_dtype{ + expected_endian_char, 'c', sizeof(std::complex)}; + EXPECT_EQ(detail::numpy_serializer::get_numpy_dtype>(), expected_dtype); + } +} + +TEST(NumPySerializerMDSpan, WriteHeader) +{ + using namespace std::string_literals; + std::ostringstream oss; + detail::numpy_serializer::header_t header{{'<', 'f', 8}, false, {2, 10, 5}}; + detail::numpy_serializer::write_header(oss, header); + EXPECT_EQ(oss.str(), + "\x93NUMPY\x01\x00"s // magic string + version (1.0) + "\x76\x00"s // HEADER_LEN = 118, in little endian + "{'descr': '{ + {"apple", "2"}, {"pie", "'is'"}, {"delicious", "True"}, {"piece of", "'cake'"}}; + EXPECT_EQ(parse, expected_parse); +} + +TEST(NumPySerializerMDSpan, ParsePyString) +{ + EXPECT_EQ(detail::numpy_serializer::parse_pystring("'foobar'"), "foobar"); +} + +TEST(NumPySerializerMDSpan, ParsePyTuple) +{ + { + std::string tuple_str{"(2,)"}; + std::vector expected_parse{"2"}; + EXPECT_EQ(detail::numpy_serializer::parse_pytuple(tuple_str), expected_parse); + } + { + std::string tuple_str{"(2, 3)"}; + std::vector expected_parse{"2", "3"}; + EXPECT_EQ(detail::numpy_serializer::parse_pytuple(tuple_str), expected_parse); + } + { + std::string tuple_str{"(2, 3, 10, 20)"}; + std::vector expected_parse{"2", "3", "10", "20"}; + EXPECT_EQ(detail::numpy_serializer::parse_pytuple(tuple_str), expected_parse); + } +} + +template +void run_roundtrip_test_scalar_serializer(T scalar) +{ + std::ostringstream oss; + detail::numpy_serializer::serialize_scalar(oss, scalar); + std::istringstream iss(oss.str()); + T tmp = detail::numpy_serializer::deserialize_scalar(iss); + EXPECT_EQ(scalar, tmp); +} + +TEST(NumPySerializerScalar, E2ERoundTrip) +{ + using namespace std::complex_literals; + run_roundtrip_test_scalar_serializer(2.0f); + run_roundtrip_test_scalar_serializer(-2.0); + run_roundtrip_test_scalar_serializer(-2); + run_roundtrip_test_scalar_serializer(0x4FFFFFF); + run_roundtrip_test_scalar_serializer>(1.0 - 2.0i); +} + +template +void check_header_scalar_serializer(T scalar) +{ + std::ostringstream oss; + detail::numpy_serializer::serialize_scalar(oss, scalar); + std::istringstream iss(oss.str()); + detail::numpy_serializer::header_t header = detail::numpy_serializer::read_header(iss); + EXPECT_TRUE(header.shape.empty()); + EXPECT_EQ(header.dtype.to_string(), detail::numpy_serializer::get_numpy_dtype().to_string()); +} + +TEST(NumPySerializerScalar, HeaderCheck) +{ + using namespace std::complex_literals; + check_header_scalar_serializer(2.0f); + check_header_scalar_serializer(-2.0); + check_header_scalar_serializer(-2); + check_header_scalar_serializer(0x4FFFFFF); + check_header_scalar_serializer>(1.0 - 2.0i); +} + +} // namespace raft diff --git a/cpp/test/neighbors/ann_ivf_flat.cu b/cpp/test/neighbors/ann_ivf_flat.cu index 8ccbe39889..98cc11c24e 100644 --- a/cpp/test/neighbors/ann_ivf_flat.cu +++ b/cpp/test/neighbors/ann_ivf_flat.cu @@ -188,10 +188,10 @@ class AnnIVFFlatTest : public ::testing::TestWithParam> { indices_ivfflat_dev.data(), ps.num_queries, ps.k); auto dists_out_view = raft::make_device_matrix_view( distances_ivfflat_dev.data(), ps.num_queries, ps.k); - raft::spatial::knn::ivf_flat::detail::save(handle_, "ivf_flat_index", index_2); + raft::spatial::knn::ivf_flat::detail::serialize(handle_, "ivf_flat_index", index_2); auto index_loaded = - raft::spatial::knn::ivf_flat::detail::load(handle_, "ivf_flat_index"); + raft::spatial::knn::ivf_flat::detail::deserialize(handle_, "ivf_flat_index"); ivf_flat::search(handle_, index_loaded, diff --git a/cpp/test/neighbors/ann_ivf_pq.cuh b/cpp/test/neighbors/ann_ivf_pq.cuh index 488041f527..31261871c1 100644 --- a/cpp/test/neighbors/ann_ivf_pq.cuh +++ b/cpp/test/neighbors/ann_ivf_pq.cuh @@ -215,9 +215,9 @@ class ivf_pq_test : public ::testing::TestWithParam { { { auto index = build_index(); - raft::spatial::knn::ivf_pq::detail::save(handle_, "ivf_pq_index", index); + raft::spatial::knn::ivf_pq::detail::serialize(handle_, "ivf_pq_index", index); } - auto index = raft::spatial::knn::ivf_pq::detail::load(handle_, "ivf_pq_index"); + auto index = raft::spatial::knn::ivf_pq::detail::deserialize(handle_, "ivf_pq_index"); size_t queries_size = ps.num_queries * ps.k; std::vector indices_ivf_pq(queries_size); diff --git a/python/pylibraft/pylibraft/common/CMakeLists.txt b/python/pylibraft/pylibraft/common/CMakeLists.txt index 3b49cef429..6ce1dfe347 100644 --- a/python/pylibraft/pylibraft/common/CMakeLists.txt +++ b/python/pylibraft/pylibraft/common/CMakeLists.txt @@ -1,5 +1,5 @@ # ============================================================================= -# Copyright (c) 2022, NVIDIA CORPORATION. +# Copyright (c) 2022-2023, 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 @@ -13,7 +13,7 @@ # ============================================================================= # Set the list of Cython files to build -set(cython_sources cuda.pyx handle.pyx interruptible.pyx) +set(cython_sources cuda.pyx handle.pyx mdspan.pyx interruptible.pyx) set(linked_libraries raft::raft) # Build all of the Cython targets diff --git a/python/pylibraft/pylibraft/common/cpp/mdspan.pxd b/python/pylibraft/pylibraft/common/cpp/mdspan.pxd index d4d0dd8f35..c3e5abb47e 100644 --- a/python/pylibraft/pylibraft/common/cpp/mdspan.pxd +++ b/python/pylibraft/pylibraft/common/cpp/mdspan.pxd @@ -1,5 +1,5 @@ # -# Copyright (c) 2022, NVIDIA CORPORATION. +# Copyright (c) 2022-2023, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -19,8 +19,9 @@ # cython: embedsignature = True # cython: language_level = 3 -import pylibraft.common.handle -from cython.operator cimport dereference as deref +from libcpp.string cimport string + +from pylibraft.common.handle cimport device_resources cdef extern from "raft/thirdparty/mdspan/include/experimental/__p0009_bits/layout_stride.hpp" namespace "std::experimental": # noqa: E501 @@ -35,6 +36,8 @@ cdef extern from "raft/core/mdspan_types.hpp" \ namespace "raft": ctypedef layout_right row_major ctypedef layout_left col_major + cdef cppclass matrix_extent[IndexType]: + pass cdef extern from "raft/core/device_mdspan.hpp" namespace "raft" nogil: @@ -73,6 +76,9 @@ cdef extern from "raft/core/host_mdspan.hpp" \ cdef cppclass host_scalar_view[ElementType, IndexType]: pass + cdef cppclass host_mdspan[ElementType, Extents, LayoutPolicy]: + pass + cdef host_matrix_view[ElementType, IndexType, LayoutPolicy] \ make_host_matrix_view[ElementType, IndexType, LayoutPolicy]( ElementType* ptr, IndexType n_rows, IndexType n_cols) except + @@ -84,3 +90,24 @@ cdef extern from "raft/core/host_mdspan.hpp" \ cdef host_scalar_view[ElementType, IndexType] \ make_host_scalar_view[ElementType, IndexType]( ElementType *ptr) except + + +cdef extern from "" namespace "std" nogil: + cdef cppclass ostringstream: + ostringstream() except + + string str() except + + + +cdef extern from "" namespace "std" nogil: + + cdef cppclass ostream: + pass + +cdef extern from "raft/core/mdspan.hpp" namespace "raft" nogil: + cdef cppclass dextents[IndentType, Rank]: + pass + +cdef extern from "raft/core/serialize.hpp" namespace "raft" nogil: + + cdef void serialize_mdspan[ElementType, Extents, LayoutPolicy]( + const device_resources& handle, ostream& os, + const host_mdspan[ElementType, Extents, LayoutPolicy]& obj) diff --git a/python/pylibraft/pylibraft/common/mdspan.pyx b/python/pylibraft/pylibraft/common/mdspan.pyx new file mode 100644 index 0000000000..ec825495f4 --- /dev/null +++ b/python/pylibraft/pylibraft/common/mdspan.pyx @@ -0,0 +1,146 @@ +# +# Copyright (c) 2023, 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. +# +# cython: profile=False +# distutils: language = c++ +# cython: embedsignature = True +# cython: language_level = 3 + +import io + +import numpy as np + +from cpython.object cimport PyObject +from cython.operator cimport dereference as deref +from libc.stddef cimport size_t +from libc.stdint cimport int32_t, int64_t, uint32_t, uint64_t, uintptr_t + +from pylibraft.common.cpp.mdspan cimport ( + col_major, + host_mdspan, + make_host_matrix_view, + matrix_extent, + ostream, + ostringstream, + row_major, + serialize_mdspan, +) +from pylibraft.common.handle cimport device_resources + +from pylibraft.common import DeviceResources + + +cdef extern from "Python.h": + Py_buffer* PyMemoryView_GET_BUFFER(PyObject* mview) + + +def run_roundtrip_test_for_mdspan(X, fortran_order=False): + if not isinstance(X, np.ndarray) or len(X.shape) != 2: + raise ValueError("Please call this function with a NumPy array with" + "2 dimensions") + handle = DeviceResources() + cdef device_resources * handle_ = \ + handle.getHandle() + cdef ostringstream oss + if X.dtype == np.float32: + if fortran_order: + serialize_mdspan[float, matrix_extent[size_t], col_major]( + deref(handle_), + oss, + + make_host_matrix_view[float, size_t, col_major]( + PyMemoryView_GET_BUFFER( + X.data).buf, + X.shape[0], X.shape[1])) + else: + serialize_mdspan[float, matrix_extent[size_t], row_major]( + deref(handle_), + oss, + + make_host_matrix_view[float, size_t, row_major]( + PyMemoryView_GET_BUFFER( + X.data).buf, + X.shape[0], X.shape[1])) + elif X.dtype == np.float64: + if fortran_order: + serialize_mdspan[double, matrix_extent[size_t], col_major]( + deref(handle_), + oss, + + make_host_matrix_view[double, size_t, col_major]( + PyMemoryView_GET_BUFFER( + X.data).buf, + X.shape[0], X.shape[1])) + else: + serialize_mdspan[double, matrix_extent[size_t], row_major]( + deref(handle_), + oss, + + make_host_matrix_view[double, size_t, row_major]( + PyMemoryView_GET_BUFFER( + X.data).buf, + X.shape[0], X.shape[1])) + elif X.dtype == np.int32: + if fortran_order: + serialize_mdspan[int32_t, matrix_extent[size_t], col_major]( + deref(handle_), + oss, + + make_host_matrix_view[int32_t, size_t, col_major]( + PyMemoryView_GET_BUFFER( + X.data).buf, + X.shape[0], X.shape[1])) + else: + serialize_mdspan[int32_t, matrix_extent[size_t], row_major]( + deref(handle_), + oss, + + make_host_matrix_view[int32_t, size_t, row_major]( + PyMemoryView_GET_BUFFER( + X.data).buf, + X.shape[0], X.shape[1])) + elif X.dtype == np.uint32: + if fortran_order: + serialize_mdspan[uint32_t, matrix_extent[size_t], col_major]( + deref(handle_), + oss, + + make_host_matrix_view[uint32_t, size_t, col_major]( + PyMemoryView_GET_BUFFER( + X.data).buf, + X.shape[0], X.shape[1])) + else: + serialize_mdspan[uint32_t, matrix_extent[size_t], row_major]( + deref(handle_), + oss, + + make_host_matrix_view[uint32_t, size_t, row_major]( + PyMemoryView_GET_BUFFER( + X.data).buf, + X.shape[0], X.shape[1])) + else: + raise NotImplementedError() + f = io.BytesIO(oss.str()) + X2 = np.load(f) + assert np.all(X.shape == X2.shape) + assert np.all(X == X2) diff --git a/python/pylibraft/pylibraft/neighbors/ivf_pq/cpp/c_ivf_pq.pxd b/python/pylibraft/pylibraft/neighbors/ivf_pq/cpp/c_ivf_pq.pxd index 3a286868bf..c56c3e9d9b 100644 --- a/python/pylibraft/pylibraft/neighbors/ivf_pq/cpp/c_ivf_pq.pxd +++ b/python/pylibraft/pylibraft/neighbors/ivf_pq/cpp/c_ivf_pq.pxd @@ -177,10 +177,10 @@ cdef extern from "raft_runtime/neighbors/ivf_pq.hpp" \ float* distances, device_memory_resource* mr) except + - cdef void save(const device_resources& handle, - const string& filename, - const index[uint64_t]& index) except + + cdef void serialize(const device_resources& handle, + const string& filename, + const index[uint64_t]& index) except + - cdef void load(const device_resources& handle, - const string& filename, - index[uint64_t]* index) except + + cdef void deserialize(const device_resources& handle, + const string& filename, + index[uint64_t]* index) except + diff --git a/python/pylibraft/pylibraft/neighbors/ivf_pq/ivf_pq.pyx b/python/pylibraft/pylibraft/neighbors/ivf_pq/ivf_pq.pyx index 42f508c969..e7b69ddbea 100644 --- a/python/pylibraft/pylibraft/neighbors/ivf_pq/ivf_pq.pyx +++ b/python/pylibraft/pylibraft/neighbors/ivf_pq/ivf_pq.pyx @@ -795,7 +795,7 @@ def save(filename, Index index, handle=None): cdef string c_filename = filename.encode('utf-8') - c_ivf_pq.save(deref(handle_), c_filename, deref(index.index)) + c_ivf_pq.serialize(deref(handle_), c_filename, deref(index.index)) @auto_sync_handle @@ -852,7 +852,7 @@ def load(filename, handle=None): cdef string c_filename = filename.encode('utf-8') index = Index() - c_ivf_pq.load(deref(handle_), c_filename, index.index) + c_ivf_pq.deserialize(deref(handle_), c_filename, index.index) index.trained = True return index diff --git a/python/pylibraft/pylibraft/test/test_mdspan_serializer.py b/python/pylibraft/pylibraft/test/test_mdspan_serializer.py new file mode 100644 index 0000000000..412cf676d0 --- /dev/null +++ b/python/pylibraft/pylibraft/test/test_mdspan_serializer.py @@ -0,0 +1,26 @@ +# Copyright (c) 2023, 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. +# + +import numpy as np +import pytest + +from pylibraft.common.mdspan import run_roundtrip_test_for_mdspan + + +# TODO(hcho3): Set up hypothesis +@pytest.mark.parametrize("dtype", ["float32", "float64", "int32", "uint32"]) +def test_mdspan_serializer(dtype): + X = np.random.random_sample((2, 3)).astype(dtype) + run_roundtrip_test_for_mdspan(X)