diskann-wide 0.48.0

DiskANN is a fast approximate nearest neighbor search library for high dimensional data
Documentation

Wide - Cross Architecture SIMD

This crate attempts to provide (mostly) Miri-compatible, cross-platform SIMD with support for light-weight architecture dispatching.

Traits

SIMD vectors are weird types as they behave both like scalars and containers. Primary traits exposed by wide are:

  • [SIMDVector]: General trait for working with a SIMD vector, including creation and data access.

  • [SIMDMask]: Basically a SIMD boolean. Comparisons between SIMDVectors are done lanewise, with the mask containing the results for each lane. Each [SIMDVector] has an associated mask.

  • [Architecture]: SIMD instructions are architecture specific. Some server CPUs like new(ish) x86 models support AVX512, while most consumer CPUs do not yet support that instruction set extension.

    To allow compilation of single binaries that support multiple architectures, wide has taken the position that the [Architecture] is largely explicit when it comes to SIMD types.

    Generic, cross-architecture algorithms are still supported by using an [Architecture]s associated SIMD types.

A host of secondary SIMD related traits are also exported, all prefixed with SIMD. Refer to the documentation on each trait for more information.

Structs

Types implementing [SIMDMask] can take a variety of architecture specific shapes. To that end, each architecture-specific [SIMDMask] is associated with a [BitMask], where bit i is set to 1 if the corresponding lane in the full mask representation evaluates to a logic true, and 0 otherwise.

Masks can be converted to and from their corresponding [BitMask] as needed.

Safety

One source of unsafety in SIMD is the accidental use of an intrinsic that is not supported by the current runtime CPU. This is made safe in wide by using the following strategy:

  • Each [SIMDVector] and [SIMDMask] type is uniquely associated with an [Architecture].

  • Construction of a new [SIMDVector] or [SIMDMask] requires either an instance of its associated architecture, or a [SIMDVector]/[SIMDMask] of the same [Architecture].

  • [Architecture] instances can only be obtained:

    • From an instance of a [SIMDVector]/[SIMDMask] associated with that [Architecture].
    • From one of the safe constructors like [arch::dispatch] or new_checked which perform runtime checks necessary to ensure the compatibility.
    • Through an unsafe constructor, on which case all bets are off.

So an [Architecture] is needed to bootstrap the use of SIMD, but from then on, the existence of SIMD types for a given [Architecture] serve as proof-of-safety.

Special Architectures

Some [Architecture]s are special and always available to use safely:

  • [arch::Scalar]: An architecture that uses emulation via loops to implement SIMD-like operations. This architecture is safe because no special hardware intrinsics are invoked.

  • [arch::Current]: The [Architecture] that is the closest fit to the current compilation target. This is not always [arch::Scalar]. For example, if compiling for x86-64-v3, then the [arch::Current] will be [arch::x86_64::V3]. This is safe because it only uses intrinsics that are already available for the compiler to use.

    The current architecture can be obtained using with [arch::current()] or the constant [crate::ARCH].

Dev Docs

Adding a new TxN vector type.

  1. Implement the type for the backends in arch (you can usually follow and slightly modify the existing examples).

  2. Implement for Emulated for the implementations that require macro instantiation.

  3. Add the type to the [Architecture] trait.

At each step, be sure to include tests, which should be fairly straight forward.

Adding a New Implementation to an Existing Trait

Basically do steps 2-4 of the above list.

Adding a New Trait

  1. If needed, provide a reference implementation in the reference module.

  2. If it's a relatively simple op, adding a new macro in test_utils/ops.rs that invokes the reference implementation may be all that's needed.

    More complicated operations may require their own test harness (see test_utils/dot_product.rs).

    Tests should go through the utilities in test_utils::driver to ensure adequate coverage and low compile time.

  3. Implement the trait for the needed types, implementing for [Emulated], architecture-specific types, [Architecture].

Testing and Architectural Levels

By default, wide will only run tests supported by the current runtime hardware. This allows the tests to pass on a wide variety of machines during development.

However, this can mean that tests targeting architecture not supported by the runtime hardware will silently succeed.

To ensure all tests either run, or generate an error if the runtime hardware does not support a test, set the environment variable

WIDE_TEST_MIN_ARCH="all"

Various back-end specific values are supported. Note that this variable sets the minimum level of tests that are required to run. Tests for higher architecture levels will still be run if supported by the runtime hardware.

x86_64

  • x86-64-v4: Target Wide's [arch::x86_64::V4] architecture.
  • x86-64-v3: Target Wide's [arch::x86_64::V3] architecture.
  • scalar: Target the scalar architecture.