|
8 | 8 | #include <cuvs/neighbors/hnsw.hpp> |
9 | 9 | #include <raft/core/logger.hpp> |
10 | 10 |
|
| 11 | +#include <algorithm> |
11 | 12 | #include <chrono> |
| 13 | +#include <filesystem> |
| 14 | +#include <fstream> |
| 15 | +#include <iterator> |
12 | 16 | #include <memory> |
13 | 17 | #include <optional> |
14 | 18 | #include <variant> |
15 | 19 |
|
16 | 20 | #include "../common/ann_types.hpp" |
| 21 | +#include "../common/blob.hpp" |
| 22 | +#include "../common/conf.hpp" |
17 | 23 | #include "../diskann/diskann_wrapper.h" |
18 | 24 | #include "cuvs_ann_bench_utils.h" |
19 | 25 | #include <cuvs/neighbors/vamana.hpp> |
@@ -165,51 +171,35 @@ void cuvs_cagra_diskann<T, IdxT>::save(const std::string& file) const |
165 | 171 | index_of.close(); |
166 | 172 | if (!index_of) { RAFT_FAIL("Error writing output %s", file.c_str()); } |
167 | 173 |
|
168 | | - // try allocating a buffer for the dataset on host |
169 | | - try { |
170 | | - auto const* idx_ptr = cagra_build_.get_index(); |
171 | | - std::optional<raft::host_matrix<T, int64_t>> h_dataset = std::nullopt; |
172 | | - auto const& data_view = idx_ptr->dataset(); |
173 | | - if constexpr (cuvs::neighbors::is_padded_dataset_view_v<std::decay_t<decltype(data_view)>>) { |
174 | | - auto const& v = data_view; |
175 | | - auto n_rows = v.n_rows(); |
176 | | - auto dim = v.dim(); |
177 | | - auto stride = v.stride(); |
178 | | - h_dataset.emplace(raft::make_host_matrix<T, int64_t>(n_rows, dim)); |
179 | | - raft::copy_matrix(h_dataset->data_handle(), |
180 | | - dim, |
181 | | - v.view().data_handle(), |
182 | | - stride, |
183 | | - dim, |
184 | | - n_rows, |
185 | | - raft::resource::get_cuda_stream(handle_)); |
186 | | - } else { |
187 | | - RAFT_LOG_DEBUG("dataset serialization: index dataset is not device_padded_dataset_view"); |
188 | | - } |
189 | | - |
190 | | - if (h_dataset.has_value()) { |
191 | | - raft::resource::sync_stream(handle_); |
192 | | - std::string dataset_base_file = file + ".data"; |
193 | | - std::ofstream dataset_of(dataset_base_file, std::ios::out | std::ios::binary); |
194 | | - if (!dataset_of) { RAFT_FAIL("Cannot open file %s", dataset_base_file.c_str()); } |
195 | | - size_t dataset_file_offset = 0; |
196 | | - int size = static_cast<int>(cagra_build_.get_index()->size()); |
197 | | - int dim = static_cast<int>(cagra_build_.get_index()->dim()); |
198 | | - dataset_of.seekp(dataset_file_offset, dataset_of.beg); |
199 | | - dataset_of.write((char*)&size, sizeof(int)); |
200 | | - dataset_of.write((char*)&dim, sizeof(int)); |
201 | | - for (int i = 0; i < size; i++) { |
202 | | - dataset_of.write((char*)(h_dataset->data_handle() + i * h_dataset->extent(1)), |
203 | | - dim * sizeof(T)); |
204 | | - } |
205 | | - dataset_of.close(); |
206 | | - if (!dataset_of) { RAFT_FAIL("Error writing output %s", dataset_base_file.c_str()); } |
207 | | - } |
208 | | - } catch (std::bad_alloc& e) { |
209 | | - RAFT_LOG_INFO("Failed to serialize dataset"); |
210 | | - } catch (raft::logic_error& e) { |
211 | | - RAFT_LOG_INFO("Failed to serialize dataset"); |
212 | | - } |
| 174 | + // Write the rows next to the graph; diskann::Index::load() reads them from `<file>.data`. |
| 175 | + // The benchmark base file is already in the same bin format, so copy it rather than pull the |
| 176 | + // rows out of memory - this way `save()` does not care where the dataset was allocated. |
| 177 | + const auto& ds_conf = configuration::singleton().get_dataset_conf(); |
| 178 | + blob_file<T> base{ds_conf.base_file, ds_conf.subset_first_row, ds_conf.subset_size}; |
| 179 | + int size = static_cast<int>(base.rows_limit()); |
| 180 | + int dim = static_cast<int>(base.n_cols()); |
| 181 | + RAFT_EXPECTS(dim == this->dim_, "base_file dimensionality does not match the index"); |
| 182 | + |
| 183 | + size_t header_bytes = 2 * sizeof(uint32_t); |
| 184 | + size_t skip_bytes = sizeof(T) * static_cast<size_t>(base.rows_offset()) * dim; |
| 185 | + size_t copy_bytes = sizeof(T) * static_cast<size_t>(size) * dim; |
| 186 | + RAFT_EXPECTS(std::filesystem::file_size(base.path()) >= header_bytes + skip_bytes + copy_bytes, |
| 187 | + "base_file is shorter than its header claims"); |
| 188 | + |
| 189 | + std::ifstream base_in(base.path(), std::ios::in | std::ios::binary); |
| 190 | + if (!base_in) { RAFT_FAIL("Cannot open file %s", base.path().c_str()); } |
| 191 | + base_in.seekg(header_bytes + skip_bytes); |
| 192 | + |
| 193 | + std::string dataset_base_file = file + ".data"; |
| 194 | + std::ofstream dataset_of(dataset_base_file, std::ios::out | std::ios::binary); |
| 195 | + if (!dataset_of) { RAFT_FAIL("Cannot open file %s", dataset_base_file.c_str()); } |
| 196 | + dataset_of.write((char*)&size, sizeof(int)); |
| 197 | + dataset_of.write((char*)&dim, sizeof(int)); |
| 198 | + std::copy_n(std::istreambuf_iterator<char>(base_in), |
| 199 | + copy_bytes, |
| 200 | + std::ostreambuf_iterator<char>(dataset_of)); |
| 201 | + dataset_of.close(); |
| 202 | + if (!base_in || !dataset_of) { RAFT_FAIL("Error writing output %s", dataset_base_file.c_str()); } |
213 | 203 | } |
214 | 204 |
|
215 | 205 | template <typename T, typename IdxT> |
|
0 commit comments