Pad the list extents to account for interleaved storage

[achirkin]

Please refer to the discussion NVIDIA#2297

[qwertyforce]

yeah, this is better and cleaner
i searched for usages of make_list_extents, i think there is another bug
serialize_list/deserialize_list use make_list_extents and
serialize/deserialize use serialize_list/deserialize_list

cuvs/cpp/src/neighbors/ivf_flat/ivf_flat_serialize.cuh

Line 79 in d40ee2d

raft::Pow2<kIndexGroupSize>::roundUp(sizes_host(label)));

i think we can just remove round up at all, because now we do roundup inside of make_list_extents?

[achirkin]

I think it's ok to keep it as is just to minimize the diff and keep the pr laser-focused (this code is benign since rounding up is idempotent).

[qwertyforce]

as i understand, deserialize_list expects logical size, because it uses this value to create the list

cuvs/cpp/src/neighbors/ivf_list.cuh

Lines 150 to 159 in d40ee2d

	enable_if_valid_list_t<ListT> deserialize_list(const raft::resources& handle,
	std::istream& is,
	std::shared_ptr<ListT>& ld,
	const typename ListT::spec_type& store_spec,
	const typename ListT::spec_type& device_spec)
	{
	using size_type = typename ListT::size_type;
	auto size = raft::deserialize_scalar<size_type>(handle, is);
	if (size == 0) { return ld.reset(); }
	std::make_shared<ListT>(handle, device_spec, size).swap(ld);

demonstration

#include <cuvs/neighbors/ivf_flat.hpp>

#include <raft/core/device_mdarray.hpp>
#include <raft/core/resource/cuda_stream.hpp>
#include <raft/matrix/init.cuh>
#include <raft/util/pow2_utils.cuh>

#include <cstdio>
#include <iostream>
#include <optional>
#include <string>
#include <unistd.h>

int main()
{
  using namespace cuvs::neighbors::ivf_flat;

  raft::resources handle;

  constexpr int64_t base_rows = 100;
  constexpr int64_t rows      = 120;
  constexpr int64_t dim       = 4;
  constexpr int64_t rounded   = raft::Pow2<kIndexGroupSize>::roundUp(rows);

  std::cout << "config: base_rows=" << base_rows << ", training_rows=" << rows
            << ", dim=" << dim << ", kIndexGroupSize=" << kIndexGroupSize
            << ", rounded_extent=" << rounded << '\n';

  auto data = raft::make_device_matrix<float, int64_t>(handle, rows, dim);
  raft::matrix::fill(handle, data.view(), 0.0f);

  index_params params;
  params.n_lists                        = 1;
  params.metric                         = cuvs::distance::DistanceType::L2Expanded;
  params.add_data_on_build              = false;
  params.kmeans_trainset_fraction       = 1.0;
  params.adaptive_centers               = false;
  params.conservative_memory_allocation = false;

  auto all_data_view =
    raft::make_device_matrix_view<const float, int64_t>(data.data_handle(), rows, dim);
  auto empty_index = build(handle, params, all_data_view);

  auto base_data_view =
    raft::make_device_matrix_view<const float, int64_t>(data.data_handle(), base_rows, dim);
  auto index = extend(handle, base_data_view, std::nullopt, empty_index);
  raft::resource::sync_stream(handle);

  std::cout << "before serialize: index.size()=" << index.size()
            << ", list[0].get_size()=" << index.lists()[0]->get_size()
            << ", list[0].indices_capacity()=" << index.lists()[0]->indices_capacity()
            << '\n';

  std::string filename =
    "/tmp/ivf_flat_serialize_size_repro_" + std::to_string(getpid()) + ".bin";

  serialize(handle, filename, index);

  cuvs::neighbors::ivf_flat::index<float, int64_t> index_loaded(handle);
  deserialize(handle, filename, &index_loaded);
  raft::resource::sync_stream(handle);

  std::cout << "after deserialize: index_loaded.size()=" << index_loaded.size()
            << ", list[0].get_size()=" << index_loaded.lists()[0]->get_size()
            << ", list[0].indices_capacity()=" << index_loaded.lists()[0]->indices_capacity()
            << '\n';

  std::remove(filename.c_str());

  if (index_loaded.lists()[0]->get_size() != base_rows) {
    std::cerr << "FAIL: loaded logical list size changed from " << base_rows << " to "
              << index_loaded.lists()[0]->get_size() << '\n';
    return 1;
  }

  std::cout << "PASS\n";
  return 0;
}

gives

config: base_rows=100, training_rows=120, dim=4, kIndexGroupSize=32, rounded_extent=128
before serialize: index.size()=100, list[0].get_size()=100, list[0].indices_capacity()=128
after deserialize: index_loaded.size()=100, list[0].get_size()=128, list[0].indices_capacity()=128
FAIL: loaded logical list size changed from 100 to 128

idk, if we should fix it in this pr or should i create another one?

[achirkin]

You're right, another bug! The fix is minimal though, so I think it's ok to put both fixes together.

[achirkin]

Oh, looks like I've pushed to your fork by accident :) Let's finish this up in the main PR