|
| 1 | +/* |
| 2 | + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. |
| 3 | + * SPDX-License-Identifier: Apache-2.0 |
| 4 | + */ |
| 5 | +#pragma once |
| 6 | + |
| 7 | +#include <cuvs/core/cuda_fp16.hpp> |
| 8 | +#include <cuvs/core/export.hpp> |
| 9 | + |
| 10 | +#include <raft/core/detail/mdspan_numpy_serializer.hpp> |
| 11 | +#include <raft/core/error.hpp> |
| 12 | + |
| 13 | +#include <cstddef> |
| 14 | +#include <cstdint> |
| 15 | +#include <limits> |
| 16 | +#include <sstream> |
| 17 | +#include <string> |
| 18 | +#include <type_traits> |
| 19 | +#include <vector> |
| 20 | + |
| 21 | +namespace CUVS_EXPORT cuvs { |
| 22 | +namespace util { |
| 23 | + |
| 24 | +template <typename T> |
| 25 | +inline auto numpy_dtype_string() -> std::string |
| 26 | +{ |
| 27 | + if constexpr (std::is_same_v<T, half>) { |
| 28 | + return "<f2"; |
| 29 | + } else { |
| 30 | + return raft::detail::numpy_serializer::get_numpy_dtype<T>().to_string(); |
| 31 | + } |
| 32 | +} |
| 33 | + |
| 34 | +inline auto make_numpy_header_from_dtype(const std::string& dtype, const std::vector<size_t>& shape) |
| 35 | + -> std::string |
| 36 | +{ |
| 37 | + std::stringstream dict; |
| 38 | + dict << "{'descr': '" << dtype << "', 'fortran_order': False, 'shape': ("; |
| 39 | + for (size_t i = 0; i < shape.size(); ++i) { |
| 40 | + if (i != 0) { dict << ", "; } |
| 41 | + dict << shape[i]; |
| 42 | + } |
| 43 | + if (shape.size() == 1) { dict << ","; } |
| 44 | + dict << "), }"; |
| 45 | + |
| 46 | + std::string header = dict.str(); |
| 47 | + constexpr size_t preamble_size = |
| 48 | + 6 + 2 + sizeof(uint16_t); // magic string, version, v1 header length |
| 49 | + const size_t remainder = (preamble_size + header.size() + 1) % 16; |
| 50 | + if (remainder != 0) { header.append(16 - remainder, ' '); } |
| 51 | + header.push_back('\n'); |
| 52 | + |
| 53 | + RAFT_EXPECTS(header.size() <= std::numeric_limits<uint16_t>::max(), |
| 54 | + "NumPy v1 header is too large: %zu bytes", |
| 55 | + header.size()); |
| 56 | + |
| 57 | + const auto header_len = static_cast<uint16_t>(header.size()); |
| 58 | + std::string result; |
| 59 | + result.reserve(preamble_size + header.size()); |
| 60 | + result.append("\x93NUMPY", 6); |
| 61 | + result.push_back(1); |
| 62 | + result.push_back(0); |
| 63 | + result.push_back(static_cast<char>(header_len & 0xff)); |
| 64 | + result.push_back(static_cast<char>((header_len >> 8) & 0xff)); |
| 65 | + result.append(header); |
| 66 | + return result; |
| 67 | +} |
| 68 | + |
| 69 | +template <typename T> |
| 70 | +inline auto make_numpy_header_string(const std::vector<size_t>& shape) -> std::string |
| 71 | +{ |
| 72 | + if constexpr (std::is_same_v<T, half>) { |
| 73 | + return make_numpy_header_from_dtype(numpy_dtype_string<T>(), shape); |
| 74 | + } else { |
| 75 | + const auto dtype = raft::detail::numpy_serializer::get_numpy_dtype<T>(); |
| 76 | + const bool fortran_order = false; |
| 77 | + const raft::detail::numpy_serializer::header_t header = {dtype, fortran_order, shape}; |
| 78 | + |
| 79 | + std::stringstream ss; |
| 80 | + raft::detail::numpy_serializer::write_header(ss, header); |
| 81 | + return ss.str(); |
| 82 | + } |
| 83 | +} |
| 84 | + |
| 85 | +} // namespace util |
| 86 | +} // namespace CUVS_EXPORT cuvs |
0 commit comments