diff --git a/diskann-benchmark-runner/src/files.rs b/diskann-benchmark-runner/src/files.rs index 355b47010..8333a0c04 100644 --- a/diskann-benchmark-runner/src/files.rs +++ b/diskann-benchmark-runner/src/files.rs @@ -59,6 +59,12 @@ impl CheckDeserialization for InputFile { } } +impl AsRef for InputFile { + fn as_ref(&self) -> &Path { + &*self + } +} + /////////// // Tests // /////////// diff --git a/diskann-benchmark-simd/examples/simd.json b/diskann-benchmark-simd/examples/simd.json index 8798f350c..2ad9b5444 100644 --- a/diskann-benchmark-simd/examples/simd.json +++ b/diskann-benchmark-simd/examples/simd.json @@ -6,1840 +6,12 @@ "content": { "query_type": "float32", "data_type": "float32", - "arch": "x86-64-v4", - "runs": [ - { - "distance": "squared_l2", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - } - ] - } - }, - { - "type": "simd-op", - "content": { - "query_type": "float32", - "data_type": "float32", - "arch": "x86-64-v3", - "runs": [ - { - "distance": "squared_l2", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - } - ] - } - }, - { - "type": "simd-op", - "content": { - "query_type": "float32", - "data_type": "float32", - "arch": "scalar", - "runs": [ - { - "distance": "squared_l2", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - } - ] - } - }, - { - "type": "simd-op", - "content": { - "query_type": "float32", - "data_type": "float32", - "arch": "reference", - "runs": [ - { - "distance": "squared_l2", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 50, - "num_measurements": 100 - } - ] - } - }, - - { - "type": "simd-op", - "content": { - "query_type": "float16", - "data_type": "float16", - "arch": "x86-64-v4", - "runs": [ - { - "distance": "squared_l2", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - } - ] - } - }, - { - "type": "simd-op", - "content": { - "query_type": "float16", - "data_type": "float16", - "arch": "x86-64-v3", - "runs": [ - { - "distance": "squared_l2", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - } - ] - } - }, - { - "type": "simd-op", - "content": { - "query_type": "float16", - "data_type": "float16", - "arch": "scalar", - "runs": [ - { - "distance": "squared_l2", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 50, - "num_measurements": 100 - } - ] - } - }, - { - "type": "simd-op", - "content": { - "query_type": "float16", - "data_type": "float16", - "arch": "reference", - "runs": [ - { - "distance": "squared_l2", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 50, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 50, - "num_measurements": 100 - } - ] - } - }, - - - { - "type": "simd-op", - "content": { - "query_type": "uint8", - "data_type": "uint8", - "arch": "x86-64-v4", - "runs": [ - { - "distance": "squared_l2", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - } - ] - } - }, - { - "type": "simd-op", - "content": { - "query_type": "uint8", - "data_type": "uint8", - "arch": "x86-64-v3", - "runs": [ - { - "distance": "squared_l2", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - } - ] - } - }, - { - "type": "simd-op", - "content": { - "query_type": "uint8", - "data_type": "uint8", - "arch": "scalar", - "runs": [ - { - "distance": "squared_l2", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - } - ] - } - }, - { - "type": "simd-op", - "content": { - "query_type": "uint8", - "data_type": "uint8", - "arch": "reference", - "runs": [ - { - "distance": "squared_l2", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - } - ] - } - }, - - - { - "type": "simd-op", - "content": { - "query_type": "int8", - "data_type": "int8", - "arch": "x86-64-v4", - "runs": [ - { - "distance": "squared_l2", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - } - ] - } - }, - { - "type": "simd-op", - "content": { - "query_type": "int8", - "data_type": "int8", "arch": "x86-64-v3", "runs": [ - { - "distance": "squared_l2", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - } - ] - } - }, - { - "type": "simd-op", - "content": { - "query_type": "int8", - "data_type": "int8", - "arch": "scalar", - "runs": [ - { - "distance": "squared_l2", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - } - ] - } - }, - { - "type": "simd-op", - "content": { - "query_type": "int8", - "data_type": "int8", - "arch": "reference", - "runs": [ - { - "distance": "squared_l2", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "squared_l2", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "inner_product", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, { "distance": "inner_product", - "dim": 768, - "num_points": 12, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 100, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 128, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 160, - "num_points": 50, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 384, - "num_points": 24, - "loops_per_measurement": 5000, - "num_measurements": 100 - }, - { - "distance": "cosine", - "dim": 768, - "num_points": 12, + "dim": 320, + "num_points": 972, "loops_per_measurement": 5000, "num_measurements": 100 } diff --git a/diskann-benchmark/src/backend/exhaustive/mod.rs b/diskann-benchmark/src/backend/exhaustive/mod.rs index c756c3451..c0ec64a66 100644 --- a/diskann-benchmark/src/backend/exhaustive/mod.rs +++ b/diskann-benchmark/src/backend/exhaustive/mod.rs @@ -11,6 +11,7 @@ mod algos; mod minmax; +mod multi; mod product; mod spherical; @@ -20,4 +21,5 @@ pub(crate) fn register_benchmarks(benchmarks: &mut Benchmarks) { spherical::register_benchmarks(benchmarks); minmax::register_benchmarks(benchmarks); product::register_benchmarks(benchmarks); + multi::register_benchmarks(benchmarks); } diff --git a/diskann-benchmark/src/backend/exhaustive/multi.rs b/diskann-benchmark/src/backend/exhaustive/multi.rs new file mode 100644 index 000000000..01ccc03ba --- /dev/null +++ b/diskann-benchmark/src/backend/exhaustive/multi.rs @@ -0,0 +1,262 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Exhaustive (brute-force) multi-vector KNN search using Chamfer distance. +//! +//! This module computes exact K-nearest neighbors for multi-vector data by +//! evaluating Chamfer distance between each query and all data points. + +use std::io::Write; + +use diskann_benchmark_runner::{ + describeln, + dispatcher::{self, DispatchRule, FailureScore, MatchScore}, + output::Output, + utils::MicroSeconds, + Any, +}; +use diskann_quantization::multi_vector::{distance::Chamfer, Mat, Standard}; +use diskann_vector::{PureDistanceFunction, distance::InnerProduct}; +use diskann_providers::model::graph::provider::async_::inmem; +use indicatif::{ProgressBar, ProgressStyle}; +use rayon::iter::{IntoParallelIterator, ParallelIterator}; +use serde::Serialize; + +use crate::{inputs, utils::datafiles}; + +const NAME: &str = "exhaustive-multi-vector"; + +pub(super) fn register_benchmarks(benchmarks: &mut diskann_benchmark_runner::registry::Benchmarks) { + benchmarks.register::>(NAME, |object, _checkpoint, output| { + match object.run(output) { + Ok(v) => Ok(serde_json::to_value(v)?), + Err(err) => Err(err), + } + }); +} + +macro_rules! write_field { + ($f:ident, $field:tt, $fmt:literal, $($expr:tt)*) => { + writeln!($f, concat!("{:>19}: ", $fmt), $field, $($expr)*) + } +} + +///////////////////// +// MultiExhaustive // +///////////////////// + +/// Dispatcher target for exhaustive multi-vector search. +pub(super) struct MultiExhaustive<'a> { + input: &'a inputs::multi::ExhaustiveSearch, +} + +impl<'a> MultiExhaustive<'a> { + fn new(input: &'a inputs::multi::ExhaustiveSearch) -> Self { + Self { input } + } + + fn run(self, mut output: &mut dyn Output) -> anyhow::Result { + let input = self.input; + writeln!(output, "{}", input)?; + + // Load data and queries + writeln!(output, "Loading multi-vector data...")?; + let data: Vec>> = datafiles::load_multi_vectors(&input.data)?; + let num_data = data.len(); + writeln!(output, " Loaded {} data points", num_data)?; + + writeln!(output, "Loading multi-vector queries...")?; + let queries: Vec>> = datafiles::load_multi_vectors(&input.queries)?; + let num_queries = queries.len(); + writeln!(output, " Loaded {} queries", num_queries)?; + + let k = input.num_nearest_neighbors; + if k > num_data { + anyhow::bail!( + "K ({}) exceeds number of data points ({})", + k, + num_data + ); + } + + let threadpool = rayon::ThreadPoolBuilder::new() + .num_threads(input.num_threads.get()) + .build()?; + + let progress = ProgressBar::with_draw_target( + Some(num_queries as u64), + output.draw_target(), + ); + progress.set_style(ProgressStyle::with_template( + "Exhaustive search [{elapsed_precise}] {wide_bar} {percent}%", + )?); + + writeln!(output, "Running exhaustive search (K={})...", k)?; + let start = std::time::Instant::now(); + + let results: Vec> = threadpool.install(|| { + (0..num_queries) + .into_par_iter() + .map(|query_idx| { + let query_mat = queries[query_idx].as_view(); + + // Compute distances to all data points + let mut distances: Vec<(f32, u32)> = data + .iter() + .enumerate() + .map(|(doc_idx, doc)| { + let dist = Chamfer::evaluate(query_mat, doc.as_view()); + (dist, doc_idx as u32) + }) + .collect(); + + // Sort by distance (ascending - lower is more similar) + distances.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal)); + + progress.inc(1); + + // Take top-K + distances.into_iter().take(k).map(|(_, id)| id).collect() + }) + .collect() + }); + + progress.finish(); + let search_time: MicroSeconds = start.elapsed().into(); + + // Write results to binary groundtruth file + writeln!(output, "Writing groundtruth to {}...", input.output)?; + write_groundtruth(&input.output, &results, k)?; + + let result = Results { + num_data, + num_queries, + k, + num_threads: input.num_threads.get(), + search_time, + qps: (num_queries as f64) / search_time.as_seconds(), + }; + + writeln!(output, "\n{}", result)?; + Ok(result) + } +} + +/// Write groundtruth results in binary format. +/// +/// Format: +/// - num_queries (u32) +/// - k (u32) +/// - For each query: k neighbor IDs (u32) +fn write_groundtruth(path: &str, results: &[Vec], k: usize) -> anyhow::Result<()> { + use std::io::BufWriter; + + let file = std::fs::File::create(path)?; + let mut writer = BufWriter::new(file); + + // Write header + writer.write_all(&(results.len() as u32).to_le_bytes())?; + writer.write_all(&(k as u32).to_le_bytes())?; + + // Write results + for neighbors in results { + for &id in neighbors.iter().take(k) { + writer.write_all(&id.to_le_bytes())?; + } + // Pad with zeros if needed + for _ in neighbors.len()..k { + writer.write_all(&0u32.to_le_bytes())?; + } + } + + writer.flush()?; + Ok(()) +} + +////////////// +// Dispatch // +////////////// + +impl dispatcher::Map for MultiExhaustive<'static> { + type Type<'a> = MultiExhaustive<'a>; +} + +impl<'a> DispatchRule<&'a inputs::multi::ExhaustiveSearch> for MultiExhaustive<'a> { + type Error = std::convert::Infallible; + + fn try_match(_from: &&'a inputs::multi::ExhaustiveSearch) -> Result { + Ok(MatchScore(0)) + } + + fn convert(from: &'a inputs::multi::ExhaustiveSearch) -> Result { + Ok(Self::new(from)) + } + + fn description( + f: &mut std::fmt::Formatter<'_>, + from: Option<&&'a inputs::multi::ExhaustiveSearch>, + ) -> std::fmt::Result { + match from { + None => { + describeln!(f, "- Exhaustive KNN search for multi-vector data")?; + describeln!(f, "- Uses Chamfer (asymmetric max-sim) distance")?; + describeln!(f, "- Outputs binary groundtruth file")?; + } + Some(_) => { + // No additional constraints to check + } + } + Ok(()) + } +} + +impl<'a> DispatchRule<&'a Any> for MultiExhaustive<'a> { + type Error = anyhow::Error; + + fn try_match(from: &&'a Any) -> Result { + from.try_match::() + } + + fn convert(from: &'a Any) -> Result { + from.convert::() + } + + fn description( + f: &mut std::fmt::Formatter<'_>, + from: Option<&&'a Any>, + ) -> std::fmt::Result { + Any::description::( + f, + from, + inputs::multi::ExhaustiveSearch::tag(), + ) + } +} + +///////////// +// Results // +///////////// + +#[derive(Debug, Serialize)] +struct Results { + num_data: usize, + num_queries: usize, + k: usize, + num_threads: usize, + search_time: MicroSeconds, + qps: f64, +} + +impl std::fmt::Display for Results { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write_field!(f, "Data Points", "{}", self.num_data)?; + write_field!(f, "Queries", "{}", self.num_queries)?; + write_field!(f, "K", "{}", self.k)?; + write_field!(f, "Threads", "{}", self.num_threads)?; + write_field!(f, "Search Time", "{:.3}s", self.search_time.as_seconds())?; + write_field!(f, "QPS", "{:.2}", self.qps)?; + Ok(()) + } +} diff --git a/diskann-benchmark/src/backend/index/benchmarks.rs b/diskann-benchmark/src/backend/index/benchmarks.rs index f8b128c4d..a8310ea18 100644 --- a/diskann-benchmark/src/backend/index/benchmarks.rs +++ b/diskann-benchmark/src/backend/index/benchmarks.rs @@ -36,7 +36,7 @@ use serde::Serialize; use super::{ build::{self, load_index, save_index, single_or_multi_insert, BuildStats}, - product, scalar, search, spherical, + multi, product, scalar, search, spherical, }; use crate::{ backend::index::{ @@ -84,53 +84,54 @@ macro_rules! register_streaming { pub(super) use register; pub(super) fn register_benchmarks(benchmarks: &mut diskann_benchmark_runner::registry::Benchmarks) { - // Full Precision - register!( - benchmarks, - "async-full-precision-f32", - FullPrecision<'static, f32> - ); - register!( - benchmarks, - "async-full-precision-f16", - FullPrecision<'static, f16> - ); - register!( - benchmarks, - "async-full-precision-u8", - FullPrecision<'static, u8> - ); - register!( - benchmarks, - "async-full-precision-i8", - FullPrecision<'static, i8> - ); - - // Dynamic Full Precision - register_streaming!( - benchmarks, - "async-dynamic-full-precision-f32", - DynamicFullPrecision<'static, f32> - ); - register_streaming!( - benchmarks, - "async-dynamic-full-precision-f16", - DynamicFullPrecision<'static, f16> - ); - register_streaming!( - benchmarks, - "async-dynamic-full-precision-u8", - DynamicFullPrecision<'static, u8> - ); - register_streaming!( - benchmarks, - "async-dynamic-full-precision-i8", - DynamicFullPrecision<'static, i8> - ); + // // Full Precision + // register!( + // benchmarks, + // "async-full-precision-f32", + // FullPrecision<'static, f32> + // ); + // register!( + // benchmarks, + // "async-full-precision-f16", + // FullPrecision<'static, f16> + // ); + // register!( + // benchmarks, + // "async-full-precision-u8", + // FullPrecision<'static, u8> + // ); + // register!( + // benchmarks, + // "async-full-precision-i8", + // FullPrecision<'static, i8> + // ); + + // // Dynamic Full Precision + // register_streaming!( + // benchmarks, + // "async-dynamic-full-precision-f32", + // DynamicFullPrecision<'static, f32> + // ); + // register_streaming!( + // benchmarks, + // "async-dynamic-full-precision-f16", + // DynamicFullPrecision<'static, f16> + // ); + // register_streaming!( + // benchmarks, + // "async-dynamic-full-precision-u8", + // DynamicFullPrecision<'static, u8> + // ); + // register_streaming!( + // benchmarks, + // "async-dynamic-full-precision-i8", + // DynamicFullPrecision<'static, i8> + // ); product::register_benchmarks(benchmarks); scalar::register_benchmarks(benchmarks); spherical::register_benchmarks(benchmarks); + multi::register_benchmarks(benchmarks); } ////////////// diff --git a/diskann-benchmark/src/backend/index/mod.rs b/diskann-benchmark/src/backend/index/mod.rs index 269887c6d..3d7b7e1ba 100644 --- a/diskann-benchmark/src/backend/index/mod.rs +++ b/diskann-benchmark/src/backend/index/mod.rs @@ -10,11 +10,15 @@ mod streaming; mod benchmarks; mod result; +// Multi-Vector +mod multi; + // Feature based backends. mod product; mod scalar; mod spherical; + pub(crate) fn register_benchmarks(benchmarks: &mut diskann_benchmark_runner::registry::Benchmarks) { benchmarks::register_benchmarks(benchmarks) } diff --git a/diskann-benchmark/src/backend/index/multi.rs b/diskann-benchmark/src/backend/index/multi.rs new file mode 100644 index 000000000..ba69e0e0d --- /dev/null +++ b/diskann-benchmark/src/backend/index/multi.rs @@ -0,0 +1,552 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::io::Write; +use std::sync::Arc; + +use diskann::{ANNResult, graph::{self, glue, DiskANNIndex}, provider}; +use diskann_benchmark_core::{self as benchmark_core, build::ids::ToId}; +use diskann_benchmark_runner::{ + self as runner, + dispatcher::{self, DispatchRule, FailureScore, MatchScore}, + registry::Benchmarks, + utils::{percentiles, MicroSeconds}, + Any, +}; +use diskann_providers::model::graph::provider::async_::{common as inmem_common, inmem}; +use diskann_quantization::multi_vector::{Mat, Standard}; +use diskann_utils::future::AsyncFriendly; + +use crate::{ + backend::index::{result::SearchResults, search::knn}, + inputs, + utils::datafiles, +}; + +//////////////////////////// +// Benchmark Registration // +//////////////////////////// + +pub(super) fn register_benchmarks(benchmarks: &mut Benchmarks) { + benchmarks.register::>( + "multi-vector-build", + |job, _checkpoint, output| { + job.run(output)?; + Ok(serde_json::to_value(())?) + }, + ); +} + +////////////// +// Dispatch // +////////////// + +/// The dispatcher target for multi-vector build operations. +pub(super) struct MultiVectorBuild<'a> { + input: &'a inputs::multi::BuildAndSearch, +} + +impl<'a> MultiVectorBuild<'a> { + fn new(input: &'a inputs::multi::BuildAndSearch) -> Self { + Self { input } + } +} + +impl dispatcher::Map for MultiVectorBuild<'static> { + type Type<'a> = MultiVectorBuild<'a>; +} + +/// Dispatch from the concrete input type. +impl<'a> DispatchRule<&'a inputs::multi::BuildAndSearch> for MultiVectorBuild<'a> { + type Error = std::convert::Infallible; + + fn try_match(_from: &&'a inputs::multi::BuildAndSearch) -> Result { + Ok(MatchScore(0)) + } + + fn convert(from: &'a inputs::multi::BuildAndSearch) -> Result { + Ok(Self::new(from)) + } + + fn description( + f: &mut std::fmt::Formatter<'_>, + _from: Option<&&'a inputs::multi::BuildAndSearch>, + ) -> std::fmt::Result { + writeln!(f, "tag: \"{}\"", inputs::multi::BuildAndSearch::tag()) + } +} + +/// Central dispatch mapping. +impl<'a> DispatchRule<&'a Any> for MultiVectorBuild<'a> { + type Error = anyhow::Error; + + fn try_match(from: &&'a Any) -> Result { + from.try_match::() + } + + fn convert(from: &'a Any) -> Result { + from.convert::() + } + + fn description(f: &mut std::fmt::Formatter<'_>, from: Option<&&'a Any>) -> std::fmt::Result { + Any::description::(f, from, inputs::multi::BuildAndSearch::tag()) + } +} + +impl<'a> MultiVectorBuild<'a> { + fn run(self, mut output: &mut dyn runner::Output) -> anyhow::Result<()> { + writeln!(output, "{}", self.input)?; + run(&self.input, output) + } +} + +//////////////// +// End-to-end // +//////////////// + +fn run( + input: &inputs::multi::BuildAndSearch, + mut output: &mut dyn runner::Output, +) -> anyhow::Result<()> { + let index = build(&input.build, output)?; + let queries: Arc<[_]> = datafiles::load_multi_vectors::(&input.search.queries)?.into(); + + // TODO: Placeholder. + let gt = datafiles::load_groundtruth(datafiles::BinFile(&input.search.groundtruth))?; + + let searcher = MultiKNN::new( + index.clone(), + queries, + benchmark_core::search::graph::Strategy::broadcast(inmem::multi::Strategy::new()), + )?; + + let steps = knn::SearchSteps::new( + input.search.reps, + &input.search.num_threads, + &input.search.runs, + ); + + let summaries = knn::run( + &searcher, + >, + steps, + )?; + + let results: Vec<_> = summaries.into_iter().map(|r| SearchResults::from(r)).collect(); + writeln!(output, "{}", crate::utils::DisplayWrapper(&*results))?; + + Ok(()) +} + +/////////// +// Build // +/////////// + +fn build( + input: &inputs::multi::Build, + output: &mut dyn runner::Output, +) -> anyhow::Result>>> +{ + let data = datafiles::load_multi_vectors::(&input.data)?; + + let dim = data.first().unwrap().vector_dim(); + + let provider = inmem::DefaultProvider::<_, _, _, provider::DefaultContext>::new_empty( + input.inmem_parameters(data.len(), dim), + inmem::multi::Precursor::::new(dim), + inmem_common::NoStore, + inmem_common::NoDeletes, + ).unwrap(); + + let index = Arc::new(graph::DiskANNIndex::new( + input.as_config().build()?, + provider, + None + )); + + let builder = Insert::new( + index.clone(), + data.into(), + inmem::multi::Strategy::new(), + benchmark_core::build::ids::Identity::new(), + ); + + let rt = benchmark_core::tokio::runtime(input.num_threads.get())?; + let _ = benchmark_core::build::build_tracked( + builder, + benchmark_core::build::Parallelism::dynamic( + diskann::utils::ONE, + input.num_threads, + ), + &rt, + Some(&super::build::ProgressMeter::new(output)), + )?; + + Ok(index) +} + +type MultiVec = Mat>; + +#[derive(Debug)] +pub struct Insert +where + DP: provider::DataProvider, + T: Copy, +{ + index: Arc>, + data: Arc<[MultiVec]>, + strategy: S, + to_id: Box>, +} + +impl Insert +where + DP: provider::DataProvider, + T: Copy, +{ + pub fn new( + index: Arc>, + data: Arc<[MultiVec]>, + strategy: S, + to_id: I, + ) -> Arc + where + I: ToId, + { + Arc::new(Self { + index, + data, + strategy, + to_id: Box::new(to_id), + }) + } +} + +impl benchmark_core::build::Build for Insert +where + DP: provider::DataProvider + provider::SetElement>, + S: glue::InsertStrategy> + Clone + AsyncFriendly, + T: AsyncFriendly + Copy, +{ + type Output = (); + + fn num_data(&self) -> usize { + self.data.len() + } + + async fn build(&self, range: std::ops::Range) -> ANNResult { + for i in range { + let context = DP::Context::default(); + self.index + .insert( + self.strategy.clone(), + &context, + &self.to_id.to_id(i)?, + &self.data[i], + ) + .await?; + } + Ok(()) + } +} + +//////////// +// Search // +//////////// + +/// A built-in helper for benchmarking K-nearest neighbors search on multi-vector indices. +/// +/// This is analogous to [`diskann_benchmark_core::search::graph::KNN`] but handles +/// multi-vector queries where each query is a `MultiVec` rather than a slice `[T]`. +#[derive(Debug)] +pub struct MultiKNN +where + DP: provider::DataProvider, + T: Copy, +{ + index: Arc>, + queries: Arc<[MultiVec]>, + strategy: benchmark_core::search::graph::Strategy, +} + +impl MultiKNN +where + DP: provider::DataProvider, + T: Copy, +{ + /// Construct a new [`MultiKNN`] searcher. + /// + /// If `strategy` is one of the container variants of [`Strategy`], its length + /// must match the number of queries. If this is the case, then the strategies + /// will have a querywise correspondence with the query collection. + /// + /// # Errors + /// + /// Returns an error if the number of elements in `strategy` is not compatible with + /// the number of queries. + pub fn new( + index: Arc>, + queries: Arc<[MultiVec]>, + strategy: benchmark_core::search::graph::Strategy, + ) -> anyhow::Result> { + strategy.length_compatible(queries.len())?; + + Ok(Arc::new(Self { + index, + queries, + strategy, + })) + } +} + +/// Additional metrics collected during [`MultiKNN`] search. +#[derive(Debug, Clone, Copy)] +#[non_exhaustive] +pub struct SearchMetrics { + /// The number of distance comparisons performed during search. + pub comparisons: u32, + /// The number of candidates expanded during search. + pub hops: u32, +} + +impl benchmark_core::search::Search for MultiKNN +where + DP: provider::DataProvider, + S: glue::SearchStrategy, DP::ExternalId> + Clone + AsyncFriendly, + T: AsyncFriendly + Copy, +{ + type Id = DP::ExternalId; + type Parameters = graph::SearchParams; + type Output = SearchMetrics; + + fn num_queries(&self) -> usize { + self.queries.len() + } + + fn id_count(&self, parameters: &Self::Parameters) -> benchmark_core::search::IdCount { + benchmark_core::search::IdCount::Fixed( + std::num::NonZeroUsize::new(parameters.k_value).unwrap_or(diskann::utils::ONE), + ) + } + + async fn search( + &self, + parameters: &Self::Parameters, + buffer: &mut O, + index: usize, + ) -> ANNResult + where + O: graph::SearchOutputBuffer + Send, + { + let context = DP::Context::default(); + let stats = self + .index + .search( + self.strategy.get(index)?, + &context, + &self.queries[index], + parameters, + buffer, + ) + .await?; + + Ok(SearchMetrics { + comparisons: stats.cmps, + hops: stats.hops, + }) + } +} + +/// An aggregated summary of multiple [`MultiKNN`] search runs. +#[derive(Debug, Clone)] +#[non_exhaustive] +pub struct Summary { + /// The [`benchmark_core::search::Setup`] used for the batch of runs. + pub setup: benchmark_core::search::Setup, + + /// The [`graph::SearchParams`] used for the batch of runs. + pub parameters: graph::SearchParams, + + /// The end-to-end latency for each repetition in the batch. + pub end_to_end_latencies: Vec, + + /// The average latency for individual queries. + pub mean_latencies: Vec, + + /// The 90th percentile latency for individual queries. + pub p90_latencies: Vec, + + /// The 99th percentile latency for individual queries. + pub p99_latencies: Vec, + + /// The recall metrics for search. + pub recall: benchmark_core::recall::RecallMetrics, + + /// The average number of distance comparisons per query. + pub mean_cmps: f64, + + /// The average number of neighbor hops per query. + pub mean_hops: f64, +} + +/// A [`benchmark_core::search::Aggregate`] for collecting the results of multiple +/// [`MultiKNN`] search runs. +pub struct Aggregator<'a, I> { + groundtruth: &'a dyn benchmark_core::recall::Rows, + recall_k: usize, + recall_n: usize, +} + +impl<'a, I> Aggregator<'a, I> { + /// Construct a new [`Aggregator`] using `groundtruth` for recall computation. + pub fn new( + groundtruth: &'a dyn benchmark_core::recall::Rows, + recall_k: usize, + recall_n: usize, + ) -> Self { + Self { + groundtruth, + recall_k, + recall_n, + } + } +} + +impl benchmark_core::search::Aggregate for Aggregator<'_, I> +where + I: benchmark_core::recall::RecallCompatible, +{ + type Output = Summary; + + fn aggregate( + &mut self, + run: benchmark_core::search::Run, + mut results: Vec>, + ) -> anyhow::Result { + let recall = match results.first() { + Some(first) => benchmark_core::recall::knn( + self.groundtruth, + None, + first.ids().as_rows(), + self.recall_k, + self.recall_n, + true, + )?, + None => anyhow::bail!("Results must be non-empty"), + }; + + let mut mean_latencies = Vec::with_capacity(results.len()); + let mut p90_latencies = Vec::with_capacity(results.len()); + let mut p99_latencies = Vec::with_capacity(results.len()); + + results.iter_mut().for_each(|r| { + match percentiles::compute_percentiles(r.latencies_mut()) { + Ok(values) => { + let percentiles::Percentiles { mean, p90, p99, .. } = values; + mean_latencies.push(mean); + p90_latencies.push(p90); + p99_latencies.push(p99); + } + Err(_) => { + let zero = MicroSeconds::new(0); + mean_latencies.push(0.0); + p90_latencies.push(zero); + p99_latencies.push(zero); + } + } + }); + + // Compute average comparisons and hops + let (total_cmps, total_hops, count) = results + .iter() + .flat_map(|r| r.output().iter()) + .fold((0u64, 0u64, 0usize), |(cmps, hops, n), o| { + (cmps + o.comparisons as u64, hops + o.hops as u64, n + 1) + }); + + let mean_cmps = if count > 0 { total_cmps as f64 / count as f64 } else { 0.0 }; + let mean_hops = if count > 0 { total_hops as f64 / count as f64 } else { 0.0 }; + + Ok(Summary { + setup: run.setup().clone(), + parameters: *run.parameters(), + end_to_end_latencies: results.iter().map(|r| r.end_to_end_latency()).collect(), + recall, + mean_latencies, + p90_latencies, + p99_latencies, + mean_cmps, + mean_hops, + }) + } +} + +///////////////// +// Knn Adapter // +///////////////// + +impl knn::Knn for Arc> +where + DP: provider::DataProvider, + MultiKNN: benchmark_core::search::Search< + Id = DP::ExternalId, + Parameters = graph::SearchParams, + Output = SearchMetrics, + >, + T: Copy, +{ + fn search_all( + &self, + parameters: Vec>, + groundtruth: &dyn benchmark_core::recall::Rows, + recall_k: usize, + recall_n: usize, + ) -> anyhow::Result> { + let results = benchmark_core::search::search_all( + self.clone(), + parameters.into_iter(), + Aggregator::new(groundtruth, recall_k, recall_n), + )?; + + Ok(results.into_iter().map(SearchResults::from).collect()) + } +} + +impl From for SearchResults { + fn from(summary: Summary) -> Self { + let Summary { + setup, + parameters, + end_to_end_latencies, + mean_latencies, + p90_latencies, + p99_latencies, + recall, + mean_cmps, + mean_hops, + .. + } = summary; + + let qps = end_to_end_latencies + .iter() + .map(|latency| recall.num_queries as f64 / latency.as_seconds()) + .collect(); + + Self { + num_tasks: setup.tasks.into(), + search_n: parameters.k_value, + search_l: parameters.l_value, + qps, + search_latencies: end_to_end_latencies, + mean_latencies, + p90_latencies, + p99_latencies, + recall: (&recall).into(), + mean_cmps: mean_cmps as f32, + mean_hops: mean_hops as f32, + } + } +} + diff --git a/diskann-benchmark/src/inputs/async_.rs b/diskann-benchmark/src/inputs/async_.rs index e12c26419..2799618d1 100644 --- a/diskann-benchmark/src/inputs/async_.rs +++ b/diskann-benchmark/src/inputs/async_.rs @@ -632,7 +632,7 @@ pub enum IndexSource { } impl CheckDeserialization for IndexSource { - fn check_deserialization(&mut self, checker: &mut Checker) -> Result<(), anyhow::Error> { + fn check_deserialization(&mut self, checker: &mut Checker) -> Result<(), anyhow::Error> { match self { IndexSource::Load(load) => load.check_deserialization(checker), IndexSource::Build(build) => build.check_deserialization(checker), diff --git a/diskann-benchmark/src/inputs/mod.rs b/diskann-benchmark/src/inputs/mod.rs index a0ae1a982..2e1cb753d 100644 --- a/diskann-benchmark/src/inputs/mod.rs +++ b/diskann-benchmark/src/inputs/mod.rs @@ -4,6 +4,7 @@ */ pub(crate) mod async_; +pub(crate) mod multi; pub(crate) mod disk; pub(crate) mod exhaustive; pub(crate) mod filters; @@ -16,6 +17,7 @@ pub(crate) fn register_inputs( exhaustive::register_inputs(registry)?; disk::register_inputs(registry)?; filters::register_inputs(registry)?; + multi::register_inputs(registry)?; Ok(()) } diff --git a/diskann-benchmark/src/inputs/multi.rs b/diskann-benchmark/src/inputs/multi.rs new file mode 100644 index 000000000..6c24052df --- /dev/null +++ b/diskann-benchmark/src/inputs/multi.rs @@ -0,0 +1,250 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::num::NonZeroUsize; + +use serde::{Deserialize, Serialize}; + +use diskann::graph::{config, StartPointStrategy}; +use diskann_benchmark_runner::{files::InputFile, utils::datatype::DataType, CheckDeserialization, Checker}; +use diskann_providers::model::graph::provider::async_::inmem::DefaultProviderParameters; + +use super::async_::{StartPointStrategyRef, TopkSearchPhase}; +use crate::inputs::{as_input, Input, Example}; + +const METRIC: diskann_vector::distance::Metric = diskann_vector::distance::Metric::InnerProduct; + +as_input!(BuildAndSearch); +as_input!(ExhaustiveSearch); + +pub(super) fn register_inputs( + registry: &mut diskann_benchmark_runner::registry::Inputs, +) -> anyhow::Result<()> { + registry.register(Input::::new())?; + registry.register(Input::::new())?; + Ok(()) +} + +#[derive(Debug, Serialize, Deserialize)] +pub(crate) struct Build { + pub(crate) data_type: DataType, + pub(crate) data: InputFile, + pub(crate) pruned_degree: usize, + pub(crate) max_degree: usize, + pub(crate) l_build: usize, + #[serde(with = "StartPointStrategyRef")] + pub(crate) start_point_strategy: StartPointStrategy, + pub(crate) alpha: f32, + pub(crate) backedge_ratio: f32, + pub(crate) num_threads: NonZeroUsize, +} + +impl Build { + pub(crate) fn as_config(&self) -> config::Builder { + config::Builder::new_with( + self.pruned_degree, + config::MaxDegree::new(self.max_degree), + self.l_build, + (METRIC).into(), + |b| { + b + .alpha(self.alpha) + .backedge_ratio(self.backedge_ratio); + } + ) + } + + pub(crate) fn inmem_parameters( + &self, + num_points: usize, + dim: usize, + ) -> DefaultProviderParameters { + DefaultProviderParameters { + max_points: num_points, + frozen_points: NonZeroUsize::new(self.start_point_strategy.count()).unwrap(), + metric: METRIC, + dim, + max_degree: self.max_degree as u32, + prefetch_lookahead: None, + prefetch_cache_line_level: None, + } + } +} + +#[derive(Debug, Serialize, Deserialize)] +pub(crate) struct BuildAndSearch { + pub(crate) build: Build, + pub(crate) search: TopkSearchPhase, +} + +// Constant for aligning field descriptions in Display implementations. +const PRINT_WIDTH: usize = 18; + +macro_rules! write_field { + ($f:ident, $field:tt, $($expr:tt)*) => { + writeln!($f, "{:>PRINT_WIDTH$}: {}", $field, $($expr)*) + } +} + +impl Build { + pub(crate) const fn tag() -> &'static str { + "multi-vector-build" + } +} + +impl Example for Build { + fn example() -> Self { + Self { + data_type: DataType::Float32, + data: InputFile::new("path/to/data"), + pruned_degree: 32, + max_degree: 64, + l_build: 50, + start_point_strategy: StartPointStrategy::Medoid, + alpha: 1.2, + backedge_ratio: 1.0, + num_threads: diskann::utils::ONE, + } + } +} + +impl CheckDeserialization for Build { + fn check_deserialization(&mut self, checker: &mut Checker) -> Result<(), anyhow::Error> { + self.data.check_deserialization(checker)?; + if self.pruned_degree > self.max_degree { + anyhow::bail!( + "Pruned degree ({}) must be less than max degree ({})", + self.pruned_degree, + self.max_degree, + ); + } + + Ok(()) + } +} + +impl std::fmt::Display for Build { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + writeln!(f, "Multi-Vector Index Build\n")?; + write_field!(f, "tag", Self::tag())?; + write_field!(f, "file", self.data.display())?; + write_field!(f, "data_type", self.data_type)?; + write_field!(f, "pruned degree", self.pruned_degree)?; + write_field!(f, "max degree", self.max_degree)?; + write_field!(f, "L-build", self.l_build)?; + write_field!(f, "start point strategy", self.start_point_strategy)?; + write_field!(f, "alpha", self.alpha)?; + write_field!(f, "backedge ratio", self.backedge_ratio)?; + write_field!(f, "build threads", self.num_threads)?; + Ok(()) + } +} + +impl BuildAndSearch { + pub(crate) const fn tag() -> &'static str { + "multi-vector-build-and-search" + } +} + +impl Example for BuildAndSearch { + fn example() -> Self { + Self { + build: Build::example(), + search: TopkSearchPhase::example(), + } + } +} + +impl CheckDeserialization for BuildAndSearch { + fn check_deserialization(&mut self, checker: &mut Checker) -> Result<(), anyhow::Error> { + self.build.check_deserialization(checker)?; + self.search.check_deserialization(checker)?; + Ok(()) + } +} + +impl std::fmt::Display for BuildAndSearch { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + writeln!(f, "Multi-Vector Index Build and Search\n")?; + write_field!(f, "tag", Self::tag())?; + write!(f, "{}", self.build)?; + Ok(()) + } +} + +/////////////////////// +// Exhaustive Search // +/////////////////////// + +/// Input for exhaustive (brute-force) multi-vector KNN search using Chamfer distance. +/// +/// This computes exact K-nearest neighbors for each query and writes the results +/// to a binary groundtruth file. +#[derive(Debug, Serialize, Deserialize)] +pub(crate) struct ExhaustiveSearch { + /// Path to the multi-vector dataset file. + pub(crate) data: InputFile, + /// Path to the multi-vector queries file. + pub(crate) queries: InputFile, + /// Path for the output groundtruth file (.bin format). + pub(crate) output: String, + /// Number of nearest neighbors to find per query. + pub(crate) num_nearest_neighbors: usize, + /// Number of threads for parallel search. + pub(crate) num_threads: NonZeroUsize, +} + +impl ExhaustiveSearch { + pub(crate) const fn tag() -> &'static str { + "exhaustive-multi-vector" + } +} + +impl Example for ExhaustiveSearch { + fn example() -> Self { + Self { + data: InputFile::new("path/to/data.mvbin"), + queries: InputFile::new("path/to/queries.mvbin"), + output: "groundtruth.bin".to_string(), + num_nearest_neighbors: 100, + num_threads: NonZeroUsize::new(8).unwrap(), + } + } +} + +impl CheckDeserialization for ExhaustiveSearch { + fn check_deserialization(&mut self, checker: &mut Checker) -> Result<(), anyhow::Error> { + self.data.check_deserialization(checker)?; + self.queries.check_deserialization(checker)?; + + if self.num_nearest_neighbors == 0 { + anyhow::bail!("num_nearest_neighbors must be greater than 0"); + } + + // Resolve output path relative to output directory if set + let output_path = std::path::Path::new(&self.output); + let output_filename = output_path + .file_name() + .unwrap_or_else(|| output_path.as_os_str()); + let resolved_path = checker.register_output(output_path.parent())?; + let full_path = resolved_path.join(output_filename); + self.output = full_path.to_string_lossy().to_string(); + + Ok(()) + } +} + +impl std::fmt::Display for ExhaustiveSearch { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + writeln!(f, "Exhaustive Multi-Vector Search (Chamfer Distance)\n")?; + write_field!(f, "tag", Self::tag())?; + write_field!(f, "data", self.data.display())?; + write_field!(f, "queries", self.queries.display())?; + write_field!(f, "output", self.output)?; + write_field!(f, "number of nearest neighbors", self.num_nearest_neighbors)?; + write_field!(f, "number of threads", self.num_threads)?; + Ok(()) + } +} diff --git a/diskann-benchmark/src/utils/datafiles.rs b/diskann-benchmark/src/utils/datafiles.rs index c4cf8b881..6bb8318e5 100644 --- a/diskann-benchmark/src/utils/datafiles.rs +++ b/diskann-benchmark/src/utils/datafiles.rs @@ -3,15 +3,18 @@ * Licensed under the MIT license. */ -use std::{io::Read, path::Path}; +use std::{io::{BufReader, Read}, path::Path}; use anyhow::Context; use bit_set::BitSet; use diskann::utils::IntoUsize; use diskann_benchmark_runner::utils::datatype::DataType; use diskann_providers::storage::StorageReadProvider; +use diskann_quantization::multi_vector::{Mat, MatRef, Standard}; use diskann_utils::views::Matrix; +use half::f16; use serde::{Deserialize, Serialize}; +use thiserror::Error; pub(crate) struct BinFile<'a>(pub(crate) &'a Path); @@ -170,3 +173,326 @@ impl From for BitSet { BitSet::from_bytes(&val.0) } } + +////////////////// +// Multi Vector // +////////////////// + +/// Errors that can occur when loading multi-vectors from a binary file. +#[derive(Debug, Error)] +pub enum MultiVectorLoadError { + #[error("IO error: {0}")] + Io(#[from] std::io::Error), + + #[error("Dimension mismatch: expected {expected}, found {found}")] + DimensionMismatch { expected: usize, found: usize }, + + #[error("Unexpected end of file while reading {context}")] + UnexpectedEof { context: &'static str }, + + #[error("Invalid header: K={k}, D={d} (both must be > 0)")] + InvalidHeader { k: u32, d: u32 }, +} + +/// Result of reading a single multi-vector: `Some((matrix, dimension))` or `None` for clean EOF. +type SingleMultiVectorResult = Result>, usize)>, MultiVectorLoadError>; + +/// Read raw f16 multi-vector data from the reader into the buffer. +/// +/// Returns `Ok(Some(mat_ref))` on success where `mat_ref` is a view over the buffer, +/// `Ok(None)` on clean EOF, or an error if the file is malformed. +fn read_multi_vector_raw<'a, R>( + reader: &mut R, + expected_dim: Option, + buffer: &'a mut Vec, +) -> Result>>, MultiVectorLoadError> +where + R: Read, +{ + // Read first byte of K - EOF here means clean end of file + let mut first_byte = [0u8; 1]; + match reader.read_exact(&mut first_byte) { + Ok(()) => {} + Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(None), + Err(e) => return Err(MultiVectorLoadError::Io(e)), + } + + // Read remaining 3 bytes of K + let mut k_rest = [0u8; 3]; + reader + .read_exact(&mut k_rest) + .map_err(|e| match e.kind() { + std::io::ErrorKind::UnexpectedEof => MultiVectorLoadError::UnexpectedEof { + context: "K header", + }, + _ => MultiVectorLoadError::Io(e), + })?; + let k = u32::from_le_bytes([first_byte[0], k_rest[0], k_rest[1], k_rest[2]]); + + // Read D (dimension) + let mut d_buf = [0u8; 4]; + reader + .read_exact(&mut d_buf) + .map_err(|e| match e.kind() { + std::io::ErrorKind::UnexpectedEof => MultiVectorLoadError::UnexpectedEof { + context: "D header", + }, + _ => MultiVectorLoadError::Io(e), + })?; + let d = u32::from_le_bytes(d_buf); + + // Validate header + if k == 0 || d == 0 { + return Err(MultiVectorLoadError::InvalidHeader { k, d }); + } + + let k = k as usize; + let d = d as usize; + + // Validate dimension consistency + if let Some(expected) = expected_dim { + if d != expected { + return Err(MultiVectorLoadError::DimensionMismatch { expected, found: d }); + } + } + + // Read K*D f16 values (reuse buffer, resizing as needed) + let num_elements = k * d; + buffer.resize(num_elements, f16::ZERO); + let byte_buf: &mut [u8] = bytemuck::must_cast_slice_mut(buffer.as_mut_slice()); + reader.read_exact(byte_buf).map_err(|e| match e.kind() { + std::io::ErrorKind::UnexpectedEof => MultiVectorLoadError::UnexpectedEof { + context: "vector data", + }, + _ => MultiVectorLoadError::Io(e), + })?; + + // Return a view over the buffer + let mat_ref = MatRef::new(Standard::new(k, d), buffer.as_slice()) + .expect("buffer size matches k * d"); + Ok(Some(mat_ref)) +} + +/// Read a single multi-vector from the reader, converting to type `T`. +/// +/// This is a thin wrapper over `read_multi_vector_raw` that handles the type conversion. +/// Only the conversion loop is monomorphized per `T`. +fn read_single_multi_vector( + reader: &mut R, + expected_dim: Option, + buffer: &mut Vec, +) -> SingleMultiVectorResult +where + T: Copy + Default + From, + R: Read, +{ + let Some(src) = read_multi_vector_raw(reader, expected_dim, buffer)? else { + return Ok(None); + }; + + // Create Mat and populate with converted data + let mut mat = Mat::new( + Standard::new(src.num_vectors(), src.vector_dim()), + T::default(), + ) + .expect("valid matrix layout"); + + for (dst_row, src_row) in mat.rows_mut().zip(src.rows()) { + for (dst, &src_val) in dst_row.iter_mut().zip(src_row.iter()) { + *dst = T::from(src_val); + } + } + + Ok(Some((mat, src.vector_dim()))) +} + +/// Load multi-vectors from a binary file. +/// +/// Each multi-vector is encoded as: +/// - K (u32): number of vectors in this multi-vector +/// - D (u32): dimension of each vector +/// - K×D f16 values: vector data in row-major order +/// +/// Returns a Vec of `Mat>` where each Mat represents one multi-vector. +/// All multi-vectors must have the same dimension D, but may have different K values. +/// +/// # Errors +/// +/// - `MultiVectorLoadError::Io`: IO error reading the file +/// - `MultiVectorLoadError::DimensionMismatch`: D values differ between multi-vectors +/// - `MultiVectorLoadError::UnexpectedEof`: File ends mid-record +/// - `MultiVectorLoadError::InvalidHeader`: K or D is zero +pub fn load_multi_vectors(path: impl AsRef) -> Result>>, MultiVectorLoadError> +where + T: Copy + Default + From, +{ + load_multi_vectors_inner::(path.as_ref()) +} + +fn load_multi_vectors_inner(path: &Path) -> Result>>, MultiVectorLoadError> +where + T: Copy + Default + From, +{ + let file = std::fs::File::open(path)?; + let mut reader = BufReader::new(file); + let mut result = Vec::new(); + let mut expected_dim: Option = None; + let mut buffer: Vec = Vec::new(); + + while let Some((mat, dim)) = read_single_multi_vector(&mut reader, expected_dim, &mut buffer)? { + expected_dim = Some(dim); + result.push(mat); + } + + Ok(result) +} + +/////////// +// Tests // +/////////// + +#[cfg(test)] +mod multi_vector_tests { + use super::*; + use std::io::Write; + use tempfile::NamedTempFile; + + fn write_multi_vector(file: &mut impl Write, k: u32, d: u32, values: &[f16]) { + file.write_all(&k.to_le_bytes()).unwrap(); + file.write_all(&d.to_le_bytes()).unwrap(); + file.write_all(bytemuck::cast_slice(values)).unwrap(); + } + + #[test] + fn test_empty_file() { + let file = NamedTempFile::new().unwrap(); + let result: Vec>> = load_multi_vectors(file.path()).unwrap(); + assert!(result.is_empty()); + } + + #[test] + fn test_single_multi_vector() { + let mut file = NamedTempFile::new().unwrap(); + // K=2 vectors, D=3 dimensions + let values: Vec = vec![ + f16::from_f32(1.0), + f16::from_f32(2.0), + f16::from_f32(3.0), + f16::from_f32(4.0), + f16::from_f32(5.0), + f16::from_f32(6.0), + ]; + write_multi_vector(&mut file, 2, 3, &values); + file.flush().unwrap(); + + let result: Vec>> = load_multi_vectors(file.path()).unwrap(); + assert_eq!(result.len(), 1); + assert_eq!(result[0].num_vectors(), 2); + assert_eq!(result[0].vector_dim(), 3); + + let row0 = result[0].get_row(0).unwrap(); + assert_eq!(row0, &[1.0, 2.0, 3.0]); + + let row1 = result[0].get_row(1).unwrap(); + assert_eq!(row1, &[4.0, 5.0, 6.0]); + } + + #[test] + fn test_multiple_multi_vectors_varying_k() { + let mut file = NamedTempFile::new().unwrap(); + + // First multi-vector: K=2, D=2 + let values1: Vec = vec![ + f16::from_f32(1.0), + f16::from_f32(2.0), + f16::from_f32(3.0), + f16::from_f32(4.0), + ]; + write_multi_vector(&mut file, 2, 2, &values1); + + // Second multi-vector: K=3, D=2 (same D, different K) + let values2: Vec = vec![ + f16::from_f32(5.0), + f16::from_f32(6.0), + f16::from_f32(7.0), + f16::from_f32(8.0), + f16::from_f32(9.0), + f16::from_f32(10.0), + ]; + write_multi_vector(&mut file, 3, 2, &values2); + file.flush().unwrap(); + + let result: Vec>> = load_multi_vectors(file.path()).unwrap(); + assert_eq!(result.len(), 2); + + assert_eq!(result[0].num_vectors(), 2); + assert_eq!(result[0].vector_dim(), 2); + + assert_eq!(result[1].num_vectors(), 3); + assert_eq!(result[1].vector_dim(), 2); + } + + #[test] + fn test_dimension_mismatch() { + let mut file = NamedTempFile::new().unwrap(); + + // First: K=1, D=2 + let values1: Vec = vec![f16::from_f32(1.0), f16::from_f32(2.0)]; + write_multi_vector(&mut file, 1, 2, &values1); + + // Second: K=1, D=3 (different D - should fail) + let values2: Vec = vec![f16::from_f32(1.0), f16::from_f32(2.0), f16::from_f32(3.0)]; + write_multi_vector(&mut file, 1, 3, &values2); + file.flush().unwrap(); + + let result: Result>>, _> = load_multi_vectors(file.path()); + assert!(matches!( + result, + Err(MultiVectorLoadError::DimensionMismatch { + expected: 2, + found: 3 + }) + )); + } + + #[test] + fn test_truncated_file() { + let mut file = NamedTempFile::new().unwrap(); + // Write only K, no D or data + file.write_all(&2u32.to_le_bytes()).unwrap(); + file.flush().unwrap(); + + let result: Result>>, _> = load_multi_vectors(file.path()); + assert!(matches!( + result, + Err(MultiVectorLoadError::UnexpectedEof { .. }) + )); + } + + #[test] + fn test_invalid_header_k_zero() { + let mut file = NamedTempFile::new().unwrap(); + write_multi_vector(&mut file, 0, 3, &[]); + file.flush().unwrap(); + + let result: Result>>, _> = load_multi_vectors(file.path()); + assert!(matches!( + result, + Err(MultiVectorLoadError::InvalidHeader { k: 0, d: 3 }) + )); + } + + #[test] + fn test_invalid_header_d_zero() { + let mut file = NamedTempFile::new().unwrap(); + write_multi_vector(&mut file, 2, 0, &[]); + file.flush().unwrap(); + + let result: Result>>, _> = load_multi_vectors(file.path()); + assert!(matches!( + result, + Err(MultiVectorLoadError::InvalidHeader { k: 2, d: 0 }) + )); + } +} + diff --git a/diskann-providers/src/model/graph/provider/async_/common.rs b/diskann-providers/src/model/graph/provider/async_/common.rs index e0262a981..9e57e9bce 100644 --- a/diskann-providers/src/model/graph/provider/async_/common.rs +++ b/diskann-providers/src/model/graph/provider/async_/common.rs @@ -19,6 +19,7 @@ use crate::{ /// Represents a range of start points for an index. /// The range includes `start` and excludes `end`. /// `start` is the first valid point, and `end - 1` is the last valid point. +#[derive(Debug, Clone, Copy)] pub struct StartPoints { start: u32, end: u32, @@ -84,6 +85,7 @@ impl Deref for VectorGuard { /// each vector starts on a cache-aligned boundary (64 byte aligned). /// To achieve this, vectors may not be densely packed into the underlying /// buffer. +#[derive(Debug)] pub struct AlignedMemoryVectorStore { store: UnsafeCell>, max_vectors: usize, @@ -421,6 +423,7 @@ impl Hybrid { pub struct Internal(pub T); #[cfg(test)] +#[derive(Debug)] pub struct TestCallCount { count: std::sync::atomic::AtomicUsize, } @@ -448,6 +451,7 @@ impl TestCallCount { } #[cfg(not(test))] +#[derive(Debug)] pub struct TestCallCount {} #[cfg(not(test))] diff --git a/diskann-providers/src/model/graph/provider/async_/inmem/mod.rs b/diskann-providers/src/model/graph/provider/async_/inmem/mod.rs index 4f45b684c..40e5ea29b 100644 --- a/diskann-providers/src/model/graph/provider/async_/inmem/mod.rs +++ b/diskann-providers/src/model/graph/provider/async_/inmem/mod.rs @@ -25,6 +25,8 @@ pub use full_precision::{ }; pub(super) use full_precision::{GetFullPrecision, Rerank}; +pub mod multi; + #[cfg(test)] pub mod product; #[cfg(test)] diff --git a/diskann-providers/src/model/graph/provider/async_/inmem/multi.rs b/diskann-providers/src/model/graph/provider/async_/inmem/multi.rs new file mode 100644 index 000000000..9a6504ade --- /dev/null +++ b/diskann-providers/src/model/graph/provider/async_/inmem/multi.rs @@ -0,0 +1,515 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::sync::RwLock; + +use bytemuck::Pod; +use diskann::{ + ANNResult, ANNError, + graph::{self, glue}, + provider, + utils::{IntoUsize, VectorRepr}, +}; +use diskann_quantization::multi_vector::{Chamfer, Mat, MatRef, Standard}; +use diskann_utils::future::AsyncFriendly; +use diskann_vector::{PureDistanceFunction, distance::Metric}; + +use crate::model::graph::provider::async_::{ + SimpleNeighborProviderAsync, common, inmem, postprocess::{self, DeletionCheck}, +}; + +type MultiVec = Mat>; +type MultiVecRef<'a, T> = MatRef<'a, Standard>; + +pub type Provider = + inmem::DefaultProvider, common::NoStore, D, provider::DefaultContext>; + +const METRIC: Metric = Metric::InnerProduct; + +// Precursor +#[derive(Debug, Clone, Copy)] +pub struct Precursor { + dim: usize, + _marker: std::marker::PhantomData, +} + +impl Precursor { + pub fn new(dim: usize) -> Self { + Self { + dim, + _marker: std::marker::PhantomData, + } + } +} + +impl common::CreateVectorStore for Precursor +where + T: Pod + AsyncFriendly, +{ + type Target = Store; + fn create( + self, + max_points: usize, + metric: Metric, + _prefetch_lookahead: Option, + ) -> Self::Target { + assert_eq!(metric, METRIC, "Only inner-product is supported"); + Store::new( + max_points, + self.dim, + ) + } +} + +#[derive(Debug)] +pub struct Store +where + T: Pod, +{ + // Guards for the fast memory store. + guards: Vec>, + pooled: common::AlignedMemoryVectorStore, + multi: Vec>>>, +} + +impl Store +where + T: Pod, +{ + fn new(max_points: usize,dim: usize) -> Self { + Self { + guards: (0..max_points).map(|_| RwLock::new(())).collect(), + pooled: common::AlignedMemoryVectorStore::with_capacity(max_points, dim), + multi: (0..max_points).map(|_| RwLock::new(None)).collect(), + } + } + + fn dim(&self) -> usize { + self.pooled.dim() + } +} + +impl common::VectorStore for Store +where + T: Pod + AsyncFriendly, +{ + fn total(&self) -> usize { + self.guards.len() + } + + fn count_for_get_vector(&self) -> usize { + 0 + } +} + +impl provider::SetElement> for Provider +where + T: VectorRepr + MeanPool, + D: AsyncFriendly, +{ + type SetError = ANNError; + type Guard = provider::NoopGuard; + + fn set_element( + &self, + _context: &provider::DefaultContext, + id: &u32, + element: &MultiVec, + ) -> impl Future> + Send { + let mut buf = vec![T::default(); self.base_vectors.dim()]; + T::mean_pool(&mut buf, element.as_view()); + + let store = &self.base_vectors; + let i = id.into_usize(); + + let _pooled_guard = store.guards[i].write().unwrap(); + let mut multi_guard = store.multi[i].write().unwrap(); + + // SAFETY: We hold the write guard for this slot. + unsafe { + store.pooled.get_mut_slice(i).copy_from_slice(&buf); + } + + *multi_guard = Some(element.clone()); + + // Success. + std::future::ready(Ok(provider::NoopGuard::new(*id))) + } +} + +#[derive(Debug)] +pub struct Accessor<'a, T, D> +where + T: Pod, +{ + provider: &'a Provider, + buffer: Vec, +} + +impl<'a, T, D> Accessor<'a, T, D> +where + T: Pod, +{ + fn new(provider: &'a Provider) -> Self { + let dim = provider.base_vectors.dim(); + Self { + provider, + buffer: vec![::zeroed(); dim], + } + } + + fn store(&self) -> &Store { + &self.provider.base_vectors + } +} + +////////////// +// Provider // +////////////// + +impl provider::HasId for Accessor<'_, T, D> +where + T: Pod, +{ + type Id = u32; +} + +impl<'a, T, D> provider::DelegateNeighbor<'a> for Accessor<'_, T, D> +where + T: Pod + AsyncFriendly, + D: AsyncFriendly, +{ + type Delegate = &'a SimpleNeighborProviderAsync; + + fn delegate_neighbor(&'a mut self) -> Self::Delegate { + self.provider.neighbors() + } +} + +/// This implementation of [`Accessor`] uses the mean-pooled versions of the vectors as +/// the primary data type for search. +/// +/// Reranking is performed using the full multi-vectors. +impl<'a, T, D> provider::Accessor for Accessor<'a, T, D> +where + T: VectorRepr, + D: AsyncFriendly, +{ + type Extended = Box<[T]>; + + /// We share a reference to the local buffer to minimize the duration of the lock. + type Element<'b> + = &'b [T] + where + Self: 'b; + + type ElementRef<'b> = &'b [T]; + + /// Choose to panic on an out-of-bounds access rather than propagate an error. + type GetError = common::Panics; + + #[inline(always)] + fn get_element( + &mut self, + id: u32, + ) -> impl Future, Self::GetError>> + Send { + // We cannot go through `Accessor::store` because we need the borrow checker to + // recognize that the borrow of the provider is disjoint from the borrow of the + // buffer. + let store = &self.provider.base_vectors; + + let id = id.into_usize(); + let _guard = match store.guards.get(id) { + Some(guard) => guard.read().unwrap(), + None => panic!("Index {} is out-of-bounds", id), + }; + + // SAFETY: We hold the associated guard for this data slot, so read access is safe. + self.buffer + .copy_from_slice(unsafe { store.pooled.get_slice(id) }); + + std::future::ready(Ok(&*self.buffer)) + } + + fn on_elements_unordered( + &mut self, + itr: Itr, + mut f: F, + ) -> impl Future> + Send + where + Self: Sync, + Itr: Iterator + Send, + F: Send + for<'b> FnMut(Self::ElementRef<'b>, Self::Id), + { + let store = self.store(); + + // We kind of just yolo it and assume that if `f` panics - we have bigger problems + // to worry about. + for id in itr { + let i = id.into_usize(); + + let _guard = store.guards[i].read().unwrap(); + f(unsafe { store.pooled.get_slice(i) }, id) + } + + std::future::ready(Ok(())) + } +} + +impl provider::BuildDistanceComputer for Accessor<'_, T, D> +where + T: VectorRepr, + D: AsyncFriendly, +{ + type DistanceComputerError = common::Panics; + type DistanceComputer = T::Distance; + + fn build_distance_computer( + &self, + ) -> Result { + Ok(T::distance(METRIC, Some(self.store().dim()))) + } +} + +impl provider::BuildQueryComputer> for Accessor<'_, T, D> +where + T: VectorRepr + MeanPool, + D: AsyncFriendly, +{ + type QueryComputerError = common::Panics; + type QueryComputer = T::QueryDistance; + + fn build_query_computer( + &self, + from: &MultiVec, + ) -> Result { + // TODO: `build_query_computer` should recieve by `&mut`. + let mut buf = vec![T::default(); self.store().dim()]; + T::mean_pool(&mut buf, from.as_view()); + Ok(T::query_distance(&buf, METRIC)) + } +} + +pub trait MeanPool: Copy + Sized { + fn mean_pool(dst: &mut [Self], x: MultiVecRef<'_, Self>); +} + +impl MeanPool for f32 { + fn mean_pool(dst: &mut [f32], x: MultiVecRef<'_, f32>) { + dst.fill(0.0); + x.rows().for_each(|r| { + dst.iter_mut().zip(r.iter()).for_each(|(d, s)| *d += *s); + }); + let scale = 1.0 / (x.num_vectors() as f32); + dst.iter_mut().for_each(|d| *d *= scale); + } +} + + +/////////////// +// Reranking // +/////////////// + +impl<'a, T, D> postprocess::AsDeletionCheck for Accessor<'a, T, D> +where + T: VectorRepr, + D: AsyncFriendly + DeletionCheck, +{ + type Checker = D; + fn as_deletion_check(&self) -> &D { + &self.provider.deleted + } +} + +#[derive(Debug, Default, Clone, Copy)] +pub struct ChamferRerank; + +impl glue::SearchPostProcess, MultiVec> for ChamferRerank +where + T: VectorRepr + MeanPool, + D: AsyncFriendly + DeletionCheck, + for<'a, 'b> Chamfer: PureDistanceFunction, MultiVecRef<'a, T>>, +{ + type Error = common::Panics; + + fn post_process( + &self, + accessor: &mut Accessor<'_, T, D>, + query: &MultiVec, + _computer: &T::QueryDistance, + candidates: I, + output: &mut B, + ) -> impl Future> + Send + where + I: Iterator>, + B: graph::SearchOutputBuffer + ?Sized, + { + let checker = &accessor.provider.deleted; + + let mut reranked: Vec<(u32, f32)> = candidates + .filter_map(|n| { + if checker.deletion_check(n.id) { + None + } else { + let multi = accessor.store().multi[n.id.into_usize()].read().unwrap(); + let v = Chamfer::evaluate(query.as_view(), multi.as_ref().unwrap().as_view()); + Some((n.id, v)) + } + }) + .collect(); + + reranked + .sort_unstable_by(|a, b| (a.1).partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)); + std::future::ready(Ok(output.extend(reranked))) + } +} + +////////// +// Glue // +////////// + +impl glue::ExpandBeam> for Accessor<'_, T, D> +where + T: VectorRepr + MeanPool, + D: AsyncFriendly, +{ +} + +impl glue::FillSet for Accessor<'_, T, D> +where + T: VectorRepr, + D: AsyncFriendly, +{ +} + +impl glue::SearchExt for Accessor<'_, T, D> +where + T: VectorRepr, + D: AsyncFriendly, +{ + fn starting_points(&self) -> impl Future>> { + std::future::ready(self.provider.starting_points()) + } +} + +//////////////// +// Strategies // +//////////////// + +#[derive(Debug, Clone, Copy)] +pub struct Strategy {} + +impl Strategy { + pub fn new() -> Self { + Self {} + } +} + +impl glue::SearchStrategy, MultiVec> for common::Internal +where + T: VectorRepr + MeanPool, + D: AsyncFriendly + DeletionCheck, +{ + type QueryComputer = T::QueryDistance; + type SearchAccessor<'a> = Accessor<'a, T, D>; + type SearchAccessorError = common::Panics; + type PostProcessor = postprocess::RemoveDeletedIdsAndCopy; + + fn search_accessor<'a>( + &'a self, + provider: &'a Provider, + _context: &'a provider::DefaultContext, + ) -> Result, Self::SearchAccessorError> { + Ok(Accessor::new(provider)) + } + + fn post_processor(&self) -> Self::PostProcessor { + Default::default() + } +} + +impl glue::SearchStrategy, MultiVec> for Strategy +where + T: VectorRepr + MeanPool, + D: AsyncFriendly + DeletionCheck, + for<'a, 'b> Chamfer: PureDistanceFunction, MultiVecRef<'a, T>>, +{ + type QueryComputer = T::QueryDistance; + type SearchAccessor<'a> = Accessor<'a, T, D>; + type SearchAccessorError = common::Panics; + type PostProcessor = glue::Pipeline; + // type PostProcessor = glue::Pipeline; + + fn search_accessor<'a>( + &'a self, + provider: &'a Provider, + _context: &'a provider::DefaultContext, + ) -> Result, Self::SearchAccessorError> { + Ok(Accessor::new(provider)) + } + + fn post_processor(&self) -> Self::PostProcessor { + Default::default() + } +} + +impl glue::PruneStrategy> for Strategy +where + T: VectorRepr, + D: AsyncFriendly, +{ + type DistanceComputer = T::Distance; + type PruneAccessor<'a> = Accessor<'a, T, D>; + type PruneAccessorError = diskann::error::Infallible; + + fn prune_accessor<'a>( + &'a self, + provider: &'a Provider, + _context: &'a provider::DefaultContext, + ) -> Result, Self::PruneAccessorError> { + Ok(Accessor::new(provider)) + } +} + +impl glue::InsertStrategy, MultiVec> for Strategy +where + T: VectorRepr + MeanPool, + D: AsyncFriendly + DeletionCheck, + for<'a, 'b> Chamfer: PureDistanceFunction, MultiVecRef<'a, T>>, +{ + type PruneStrategy = Self; + fn prune_strategy(&self) -> Self::PruneStrategy { + *self + } +} + +impl glue::InplaceDeleteStrategy> for Strategy +where + T: VectorRepr + MeanPool, + D: AsyncFriendly + DeletionCheck, +{ + type DeleteElementError = common::Panics; + type DeleteElement<'a> = MultiVec; + type DeleteElementGuard = diskann_utils::reborrow::Place>; + type PruneStrategy = Self; + type SearchStrategy = common::Internal; + fn search_strategy(&self) -> Self::SearchStrategy { + common::Internal(Self::new()) + } + + fn prune_strategy(&self) -> Self::PruneStrategy { + Self::new() + } + + async fn get_delete_element<'a>( + &'a self, + provider: &'a Provider, + _context: &'a provider::DefaultContext, + id: u32, + ) -> Result { + let id = id.into_usize(); + Ok(diskann_utils::reborrow::Place(provider.base_vectors.multi[id].read().unwrap().as_ref().unwrap().clone())) + } +} diff --git a/diskann-providers/src/model/graph/provider/async_/inmem/provider.rs b/diskann-providers/src/model/graph/provider/async_/inmem/provider.rs index 6e10e5330..4eaa67c51 100644 --- a/diskann-providers/src/model/graph/provider/async_/inmem/provider.rs +++ b/diskann-providers/src/model/graph/provider/async_/inmem/provider.rs @@ -230,6 +230,7 @@ use crate::{ /// TableBasedDeletes, /// ); /// ``` +#[derive(Debug)] pub struct DefaultProvider { /// The primary vector store that holds the main representation of vectors. pub base_vectors: U, diff --git a/diskann-providers/src/model/graph/provider/async_/simple_neighbor_provider.rs b/diskann-providers/src/model/graph/provider/async_/simple_neighbor_provider.rs index 319081449..825e838af 100644 --- a/diskann-providers/src/model/graph/provider/async_/simple_neighbor_provider.rs +++ b/diskann-providers/src/model/graph/provider/async_/simple_neighbor_provider.rs @@ -19,6 +19,7 @@ use crate::storage::{ self, AsyncIndexMetadata, AsyncQuantLoadContext, DiskGraphOnly, LoadWith, SaveWith, }; +#[derive(Debug)] pub struct SimpleNeighborProviderAsync { // Each adjacency list is stored in a fixed size slice of size max_degree * graph_slack_factor + 1. // The length of the list is stored in the extra element at the end. diff --git a/diskann-quantization/src/lib.rs b/diskann-quantization/src/lib.rs index 248d02f23..ea674d659 100644 --- a/diskann-quantization/src/lib.rs +++ b/diskann-quantization/src/lib.rs @@ -204,21 +204,21 @@ pub enum Parallelism { #[cfg(feature = "codegen")] pub mod __codegen; -// Miri can't handle `trybuild`. -#[cfg(all(test, not(miri)))] -mod tests { - #[test] - fn compile_tests() { - let t = trybuild::TestCases::new(); - // Begin with a `pass` test to force full compilation of all the test binaries. - // - // This ensures that post-monomorphization errors are tested. - t.pass("tests/compile-fail/bootstrap/bootstrap.rs"); - t.compile_fail("tests/compile-fail/*.rs"); - t.compile_fail("tests/compile-fail/error/*.rs"); - t.compile_fail("tests/compile-fail/multi/*.rs"); - } -} +// // Miri can't handle `trybuild`. +// #[cfg(all(test, not(miri)))] +// mod tests { +// #[test] +// fn compile_tests() { +// let t = trybuild::TestCases::new(); +// // Begin with a `pass` test to force full compilation of all the test binaries. +// // +// // This ensures that post-monomorphization errors are tested. +// t.pass("tests/compile-fail/bootstrap/bootstrap.rs"); +// t.compile_fail("tests/compile-fail/*.rs"); +// t.compile_fail("tests/compile-fail/error/*.rs"); +// t.compile_fail("tests/compile-fail/multi/*.rs"); +// } +// } #[cfg(test)] mod test_util; diff --git a/diskann-quantization/src/multi_vector/distance/mod.rs b/diskann-quantization/src/multi_vector/distance/mod.rs index a04c1277c..0b816428c 100644 --- a/diskann-quantization/src/multi_vector/distance/mod.rs +++ b/diskann-quantization/src/multi_vector/distance/mod.rs @@ -49,4 +49,4 @@ mod max_sim; mod simple; pub use max_sim::{Chamfer, MaxSim, MaxSimError}; -pub use simple::QueryMatRef; +pub use simple::{QueryMatRef, test_function}; diff --git a/diskann-quantization/src/multi_vector/distance/simple.rs b/diskann-quantization/src/multi_vector/distance/simple.rs index 736b5fbb2..7c66a2904 100644 --- a/diskann-quantization/src/multi_vector/distance/simple.rs +++ b/diskann-quantization/src/multi_vector/distance/simple.rs @@ -76,7 +76,7 @@ impl SimpleKernel { /// * `f` - Callback invoked with `(query_index, similarity)` for each query vector #[inline] pub(crate) fn max_sim_kernel( - query: QueryMatRef<'_, Standard>, + query: MatRef<'_, Standard>, doc: MatRef<'_, Standard>, mut f: F, ) where @@ -102,6 +102,13 @@ impl SimpleKernel { } } +pub fn test_function( + x: MatRef<'_, Standard>, + y: MatRef<'_, Standard>, +) -> f32 { + Chamfer::evaluate(x, y) +} + //////////// // MaxSim // //////////// @@ -128,7 +135,7 @@ where return Err(MaxSimError::InvalidBufferLength(size, n_queries)); } - SimpleKernel::max_sim_kernel(query, doc, |i, score| { + SimpleKernel::max_sim_kernel(*query, doc, |i, score| { // SAFETY: We asserted that self.size() == query.num_vectors(), // and i < query.num_vectors() due to the kernel loop bound. unsafe { *self.scores.get_unchecked_mut(i) = score }; @@ -142,13 +149,13 @@ where // Chamfer // ///////////// -impl PureDistanceFunction>, MatRef<'_, Standard>, f32> +impl PureDistanceFunction>, MatRef<'_, Standard>, f32> for Chamfer where InnerProduct: for<'a, 'b> PureDistanceFunction<&'a [T], &'b [T], f32>, { #[inline(always)] - fn evaluate(query: QueryMatRef<'_, Standard>, doc: MatRef<'_, Standard>) -> f32 { + fn evaluate(query: MatRef<'_, Standard>, doc: MatRef<'_, Standard>) -> f32 { let mut sum = 0.0f32; SimpleKernel::max_sim_kernel(query, doc, |_i, score| { @@ -159,6 +166,17 @@ where } } +impl PureDistanceFunction>, MatRef<'_, Standard>, f32> + for Chamfer +where + Self: for<'a, 'b> PureDistanceFunction>, MatRef<'b, Standard>, f32>, +{ + #[inline(always)] + fn evaluate(query: QueryMatRef<'_, Standard>, doc: MatRef<'_, Standard>) -> f32 { + Self::evaluate(*query, doc) + } +} + #[cfg(test)] mod tests { use super::*; @@ -271,7 +289,7 @@ mod tests { } // Check that SimpleKernel is also correct. - SimpleKernel::max_sim_kernel(query, doc, |i, score| { + SimpleKernel::max_sim_kernel(*query, doc, |i, score| { assert!((scores[i] - score).abs() <= 1e-6) }); @@ -299,7 +317,7 @@ mod tests { // No query vectors means sum is 0 assert_eq!(result, 0.0); - let result = Chamfer::evaluate(doc.into(), query.deref().reborrow()); + let result = Chamfer::evaluate(doc, query.deref().reborrow()); assert_eq!(result, 0.0); } diff --git a/diskann-quantization/src/multi_vector/matrix.rs b/diskann-quantization/src/multi_vector/matrix.rs index d834784f2..565d64716 100644 --- a/diskann-quantization/src/multi_vector/matrix.rs +++ b/diskann-quantization/src/multi_vector/matrix.rs @@ -242,6 +242,14 @@ pub unsafe trait NewOwned: ReprOwned { #[derive(Debug, Clone, Copy)] pub struct Defaulted; +/////////////// +// Auxiliary // +/////////////// + +pub trait NewCloned: ReprOwned { + fn new_cloned(v: MatRef<'_, Self>) -> Mat; +} + ////////////// // Standard // ////////////// @@ -442,6 +450,18 @@ where } } +// TODO: The `Default` bound is super unfortunate. +impl NewCloned for Standard +where + T: Copy + Default, +{ + fn new_cloned(v: MatRef<'_, Self>) -> Mat { + let mut new = Mat::new(*v.repr(), Defaulted).unwrap(); + new.rows_mut().zip(v.rows()).for_each(|(dst, src)| dst.copy_from_slice(src)); + new + } +} + ///////// // Mat // ///////// @@ -575,6 +595,15 @@ impl Drop for Mat { } } +impl Clone for Mat +where + T: NewCloned, +{ + fn clone(&self) -> Self { + T::new_cloned(self.reborrow()) + } +} + impl Mat> { /// Returns the raw dimension (columns) of the vectors in the matrix. #[inline] @@ -661,6 +690,13 @@ impl<'a, T: Repr> MatRef<'a, T> { Rows::new(*self) } + pub fn to_owned(&self) -> Mat + where + T: NewCloned, + { + T::new_cloned(*self) + } + /// Construct a new [`MatRef`] over the raw pointer and representation without performing /// any validity checks. /// diff --git a/diskann-utils/src/reborrow.rs b/diskann-utils/src/reborrow.rs index 9cf854a66..eb3e4ed4b 100644 --- a/diskann-utils/src/reborrow.rs +++ b/diskann-utils/src/reborrow.rs @@ -308,6 +308,9 @@ where /// A container for types `T` providing an implementation `Deref` as well as /// reborrowing functionality mapped to `T as Deref` and `T as DerefMut`. +/// +/// Additionally, [`Place`] can be used in situations where a type `T` needs to [`Lower`] +/// to itself. #[derive(Debug, Clone, Copy, PartialEq)] #[repr(transparent)] pub struct Place(pub T); @@ -345,6 +348,15 @@ where } } +impl<'a, T> Lower<'a, T> for Place +{ + type Proxy = &'a T; + fn lower(&'a self) -> Self::Proxy { + &self.0 + } +} + + /// A container that reborrows by cloning the contained value. /// /// Note that [`ReborrowMut`] is not implemented for this type. diff --git a/diskann/src/provider.rs b/diskann/src/provider.rs index f6b9c53ae..c8acbbe46 100644 --- a/diskann/src/provider.rs +++ b/diskann/src/provider.rs @@ -805,7 +805,7 @@ pub trait DefaultAccessor: DataProvider { /// A light-weight struct implementing [`ExecutionContext`]. /// /// Used for situations where a more refined execution context is not needed. -#[derive(Default, Clone)] +#[derive(Debug, Default, Clone)] pub struct DefaultContext; impl std::fmt::Display for DefaultContext {