[diskann-inmem] Prepare code for quantization and beyond

[Mark Hildebrand (hildebrandmw)]

Prepare diskann-inmem for quantization and beyond.

Note that this is infrastructure work to get the code ready. Quantization is not yet integrated.

Goals

The requirements to support quantization are annoyingly orthogonal:

1. Collections like PQ, scalar, and spherical quantization should be able to run in "quant-only" mode and "quant + full-precision" mode, where the full-precision store is used for reranking. Internally, I would like these to be the same type to cut down on unnecessary monomorphization. This means we need support for at least two collections managed by the same epoch protected Store (one quant, the other full-precision). In this case, the full-precision store can do without the invasive tags used for the primary store, but also needs to be optional and ideally support any of f32, f16, u8, i8 and beyond.

2. For testing purposes, we probably want to retain the ability for PQ to do hybrid pruning (part full-precision, part quantized). This completely breaks the current model used by inmem of managing raw &[u8] slices. While we could technically make it work, if we needed to do something like multi-vector operations, a &[u8] is just not the right approach anyways.

3. We also want to be able to support multi-vectors and other kinds of non-uniform data in an epoch guarded Store, which gets rid of the uniform assumption of the current invasive store.

Supporting all of these required a pretty drastic reorganization of diskann-inmem.

Architecture

This PR is all about moving things up and down. The architecture went from this:

Provider

                   Prune
          +======= ExpandBeam  ===+
          |                       |
   +------|-----------------------|------+
   |      |                       |      |
   |  +-------+      +----------------+  |
   |  | layer |      | invasive-store |  |
   |  +-------+      +----------------+  |
   |                                     |
   +-------------------------------------+

where the provider had a separate Layer and Store and combined the two to make a SearchAccessor and PruneAccessor to this

Provider ------ diskann::graph::glue ----+
   +--- drives --+                       |
                 |     SearchAccessor <--+
                 |     PruneAccessor
                 |             ^
                 V             |
+------------------------------|--------+
| Layer ---- bypses ----+      |        |
|   +--- drives --+     |      |        |
|                 |     +-- ExpandBeam  |
|                 v     |   Prune       |
|  +--------------------|-----+         |
|  | Store - drives -+  |     |         |
|  |                 |  |     |         |
|  |                 v  v     |         |
|  |               +--------+ |         |
|  |               | Plugin | |         |
|  |               +--------+ |         |
|  +--------------------------+         |
+---------------------------------------+

From top to bottom:

• The Provider now contains a Layer instead of a Layer and a Store. Instead, the Store has been moved directly into the Layer. Instead of the Provider being responsible for building a SearchAccessor from pieces exposed by the old Layer, the Provider completely delegates SearchAccessor and PruneAccessor construction to its Layer.

  The rest of the Provider's job is interfacing the diskann::graph::glue API to the simplified Layer API.

• Layer now gains the Store. The Layer family of traits is extended to include logical operations like insert and retire. As mentions above, it is also responsible for building SearchAccessors and PruneAccessors.

• Store receives minor changes - its internal Buffer where it used to manage the invasive data store directly has now been moved to a Plugin trait. In this architecture, the Store is just responsible for driving the Plugin trait's lifecycle API and is completely uninvolved with the mechanics of raw data reading and writing.

• Plugin is a new trait for a slot-based store whose slots are driven by Store. The old invasive store is an example of such a Plugin.

In addition to this hierarchical layering, a Layer is allowed (and indeed, expected) to bypass its immediate Store to the underlying Plugin directly to build the various accessors. For the Invasive store, this works because concurrency tags are embedded directly in the Plugin - readers do not need anything to do with the parent store.

Why This Mostly (Probably) Works

The Plugin trait provides an extension point for managing different types of data. For the quantized case, we reuse Invasive for the quantized data and a simpler Buffer-based one for the full-precision reranking data. Then we can create Plugin consisting of both layers. The short-cut from a Layer to its Plugin means a quantized provider can create it's SearchAccessor/PruneAccessors with knowledge of both. We could even have an Invasive store for both the quantized data and the full-precision data and reuse the existing full-precision infrastructure to support both quantized and full-precision searches over the same `Layer.

Additionally, Plugin makes no requirements on the kind of data store. This allows us to store un-even sized allocations in a Plugin. It's the Layer's job to make sense of everything.

Finally, the Layer knowing the details of its Plugins means we can (with some creativity) still support hybrid pruning.

Suggested Reviewing Order

This is a large PR, but I tried very hard to keep things structured. The reviewing order outlined here is a suggested bottom-up order. Understanding how the lower levels work is important for understanding how the higher ones come together.

• num.rs: A quick warm-up. This PR introduces some strongly typed integers with specific semantics.

• prefetch.rs: Another warm-up. I wasn't satisfied with the safety/flexibility of prefetching in the current in-mem provider. This PR exacerbated the situation, so I introduced a bit more structure on prefetchers.

• store/plugin.rs: This defines the Plugin trait and it's expected lifecycle. I captured the nuances in healthy module level and trait level documentation. This is probably the most nuanced change in this PR.

• store/checked.rs: An implementation of Plugin that aggressively checks that the invariants required for the Plugin API are upheld by Store. Again, there is healthy module-level documentation describing the logic.

• store/invasive.rs: The old invasive data store moved to implement the Plugin trait. This largely preserves what was already in the old Store and is conceptually much simpler than store/checked.rs.

• store/mod.rs: Modifying Store to work against a Plugin generic instead of directly managing the invasive store. Note that there are some changes to initialization. Store::new now takes three distinct arguments:

  • Layout: Description for capacity, number of frozen points, and maximum degree.
  • Config: Configuration state dedicated directly to management of the internal concurrency data structures.
  • PluginConfig: Configuration of the internal plugin. The XConfig traits are used in this PR to perform deferred initialization of large data structure.

• Detour: With the introduction of a modularized Store, changes were made to the following integration-test related files to enable concurrent stress-test of different Store implemntations.

  • src/integration/store/*: Shared boiler plate and implementations for exposing different Store to the integration test framework. This PR exposes wrappers for the Invasive and Checked stores in the invasive and checked modules respectively.
  • integration/store/*: Integration test exposure for the different stores as different jobs in the integration test suite.
    For these, the overall structure of the exposed stores is extremely similar. Most of churn in these files is moving things around so keep the amount of repeated code to a minimum.

• layer/mod.rs: Reworked Layer/Set/Search/Insert traits for the new architecture.

  • Layer: Gains a few life-cycle related items.
  • Set: Is now responsible for also obtaining an internal slot for the inserted element.
  • Search/Insert: Reworked to return SearchAccessor and PruneAccessors directly. This is mainly to allow us to keep the internal Store/Plugin details hidden from the public interface.

  In addition:

  • ExpandBeam: Moves from its old location in provider.rs. Otherwise, is mostly unchanged.
  • Prune: A new trait for pruning. Prune implementation now internally "buffer" items in the prune set. This effectively makes the type of the elements being pruned hidden, allowing for hybrid pruning in the future.

• layer/full.rs: This is where everything comes together. The Config is used to group together full-precision related constructor arguments and Full now gains a Store. Importantly, now that Full knows the full details of its store, we can more aggressively optimize ExpandBeam with fewer bounds and length checks. Note that start point initialization is now managed by Full's constructor.

  The implementation of ExpandBeam is taken pretty much directly from the old provider.rs code. Additionally, the tests have gotten more robust with Miri having more coverage of the ExpandBeam implementation and correctness tests for the various distance specializations.

AI Disclosure: An agent was used review changes and implementation details, help brainstorm, and make focused edits to documentation.